From 6104de1e902815462a965dd2668f096c726ec9a8 Mon Sep 17 00:00:00 2001 From: Mike Urbach Date: Thu, 14 Nov 2024 13:39:20 -0800 Subject: [PATCH 1/2] WIP - try to do a wholesale clone and cleanup. --- .../FIRRTL/Transforms/LowerClasses.cpp | 58 +++++++++++-------- 1 file changed, 35 insertions(+), 23 deletions(-) diff --git a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp index c9da6ce5c56f..8cde42c3fbb1 100644 --- a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp +++ b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp @@ -1151,12 +1151,17 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, mapping.map(inputValue, parameterValue); } - // Clone the property ops from the FIRRTL Class or Module to the OM Class. - SmallVector opsToErase; - OpBuilder builder = OpBuilder::atBlockBegin(classOp.getBodyBlock()); - llvm::SmallVector fieldLocs; - llvm::SmallVector fieldValues; - for (auto &op : moduleLike->getRegion(0).getOps()) { + // Clone over the body block from the module like to the OM class. This + // actually lands the body into the OM class as a second block, which we then + // manually splice into the entry block and remove later. + moduleLike->getRegion(0).cloneInto(&classOp->getRegion(0), + classOp->getRegion(0).end(), mapping); + classBody->getOperations().splice( + classBody->end(), classOp->getRegion(0).back().getOperations()); + + // Helper to check if an op is a property op, which should be removed in the + // module like, or kept in the OM class. + auto isPropertyOp = [&](Operation &op) { // Check if any operand is a property. auto propertyOperands = llvm::any_of(op.getOperandTypes(), [](Type type) { return isa(type); @@ -1169,28 +1174,24 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, auto propertyResults = llvm::any_of( op.getResultTypes(), [](Type type) { return isa(type); }); - // If there are no properties here, move along. - if (!needsClone && !propertyOperands && !propertyResults) - continue; + return needsClone || propertyOperands || propertyResults; + }; + + llvm::SmallVector fieldLocs; + llvm::SmallVector fieldValues; + for (Operation &op : + llvm::make_early_inc_range(classOp.getBodyBlock()->getOperations())) { + if (!isPropertyOp(op)) + op.erase(); - bool isField = false; if (auto propAssign = dyn_cast(op)) { if (isa(propAssign.getDest())) { // Store any output property assignments into fields op inputs. fieldLocs.push_back(op.getLoc()); - fieldValues.push_back(mapping.lookup(propAssign.getSrc())); - isField = true; + fieldValues.push_back(mapping.lookupOrDefault(propAssign.getSrc())); + propAssign.erase(); } } - - if (!isField) - // Clone the op over to the OM Class. - builder.clone(op, mapping); - - // In case this is a Module, remember to erase this op, unless it is an - // instance. Instances are handled later in updateInstances. - if (!isa(op)) - opsToErase.push_back(&op); } // If there is a 'containingModule', add an argument for 'ports', and a field. @@ -1201,14 +1202,25 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, fieldValues.push_back(argumentValue); } + OpBuilder builder = OpBuilder::atBlockEnd(classOp.getBodyBlock()); classOp.addNewFieldsOp(builder, fieldLocs, fieldValues); + // Now that we've finished the Class body block, clean out the block where the + // module like region was landed. + classOp->getRegion(0).back().erase(); + // If the module-like is a Class, it will be completely erased later. // Otherwise, erase just the property ports and ops. if (!isa(moduleLike.getOperation())) { // Erase ops in use before def order, thanks to FIRRTL's SSA regions. - for (auto *op : llvm::reverse(opsToErase)) - op->erase(); + for (auto &op : + llvm::make_early_inc_range(llvm::reverse(moduleLike.getOperation() + ->getRegion(0) + .getBlocks() + .front() + .getOperations()))) + if (isPropertyOp(op) && !isa(op)) + op.erase(); // Erase property typed ports. moduleLike.erasePorts(portsToErase); From 89971ce4ef02caa467b9645976458cc7d4286391 Mon Sep 17 00:00:00 2001 From: Hideto Ueno Date: Fri, 15 Nov 2024 18:39:12 +0900 Subject: [PATCH 2/2] [FIRRTL] Refactor class lowering to avoid unnecessary cloning Simplify the class lowering logic by directly moving or cloning operations instead of cloning the entire module body. This change: - Removes the need to clone the entire module region - Only clones instance operations that need to be preserved - Directly moves property operations to the OM class - Removes redundant block cleanup --- .../FIRRTL/Transforms/LowerClasses.cpp | 60 +++++++++++++------ 1 file changed, 43 insertions(+), 17 deletions(-) diff --git a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp index 8cde42c3fbb1..acf00f4fb1eb 100644 --- a/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp +++ b/lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp @@ -1021,6 +1021,7 @@ static om::ClassLike convertClass(FModuleLike moduleLike, OpBuilder builder, fieldTypes.push_back( NamedAttribute(name, TypeAttr::get(op.getSrc().getType()))); } + checkAddContainingModulePorts(hasContainingModule, builder, fieldNames, fieldTypes); return builder.create(moduleLike.getLoc(), name, @@ -1143,21 +1144,16 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, for (size_t i = 0; i < nAltBasePaths; ++i) classBody->addArgument(basePathType, unknownLoc); + // Mapping from block arguments or instance results to new values. + DenseMap valueMapping; for (auto inputProperty : inputProperties) { BlockArgument parameterValue = classBody->addArgument(inputProperty.type, inputProperty.loc); BlockArgument inputValue = moduleLike->getRegion(0).getArgument(inputProperty.index); - mapping.map(inputValue, parameterValue); - } - // Clone over the body block from the module like to the OM class. This - // actually lands the body into the OM class as a second block, which we then - // manually splice into the entry block and remove later. - moduleLike->getRegion(0).cloneInto(&classOp->getRegion(0), - classOp->getRegion(0).end(), mapping); - classBody->getOperations().splice( - classBody->end(), classOp->getRegion(0).back().getOperations()); + valueMapping[inputValue] = parameterValue; + } // Helper to check if an op is a property op, which should be removed in the // module like, or kept in the OM class. @@ -1177,18 +1173,52 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, return needsClone || propertyOperands || propertyResults; }; + // Move operations from the module like to the OM class. + // We need to clone only instances and register their results to the value + // mapping. Operands must be remapped using the value mapping. + for (auto &op : llvm::make_early_inc_range( + moduleLike->getRegion(0).front().getOperations())) { + if (!isPropertyOp(op)) + continue; + + Operation *newOp = &op; + + bool needsMove = true; + + // Clone instances and register their results to the value mapping. + if (isa(op)) { + newOp = op.clone(); + for (auto [oldResult, newResult] : + llvm::zip(op.getResults(), newOp->getResults())) + valueMapping[oldResult] = newResult; + classBody->getOperations().insert(classBody->end(), newOp); + needsMove = false; + } + + // Remap operands using the value mapping. + for (size_t i = 0; i < newOp->getNumOperands(); ++i) { + auto operand = newOp->getOperand(i); + auto it = valueMapping.find(operand); + if (it != valueMapping.end()) + newOp->setOperand(i, it->second); + } + + // Move the op into the OM class body. + if (needsMove) + newOp->moveBefore(classBody, classBody->end()); + } + llvm::SmallVector fieldLocs; llvm::SmallVector fieldValues; for (Operation &op : llvm::make_early_inc_range(classOp.getBodyBlock()->getOperations())) { - if (!isPropertyOp(op)) - op.erase(); + assert(isPropertyOp(op)); if (auto propAssign = dyn_cast(op)) { - if (isa(propAssign.getDest())) { + if (auto blockArg = dyn_cast(propAssign.getDest())) { // Store any output property assignments into fields op inputs. fieldLocs.push_back(op.getLoc()); - fieldValues.push_back(mapping.lookupOrDefault(propAssign.getSrc())); + fieldValues.push_back(propAssign.getSrc()); propAssign.erase(); } } @@ -1205,10 +1235,6 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike, OpBuilder builder = OpBuilder::atBlockEnd(classOp.getBodyBlock()); classOp.addNewFieldsOp(builder, fieldLocs, fieldValues); - // Now that we've finished the Class body block, clean out the block where the - // module like region was landed. - classOp->getRegion(0).back().erase(); - // If the module-like is a Class, it will be completely erased later. // Otherwise, erase just the property ports and ops. if (!isa(moduleLike.getOperation())) {