[LowerToHW] Lower all op bodies in parallel (#8383)

This commit is just a minor simplification to LowerToHW.

Author: rwy7

lib/Conversion/FIRRTLToHW/LowerToHW.cpp

@@ -570,12 +570,13 @@ struct FIRRTLModuleLowering
   LogicalResult
   lowerModulePortsAndMoveBody(FModuleOp oldModule, hw::HWModuleOp newModule,
                               CircuitLoweringState &loweringState);
-  LogicalResult lowerModuleOperations(hw::HWModuleOp module,
-                                      CircuitLoweringState &loweringState);
+  LogicalResult lowerModuleBody(hw::HWModuleOp module,
+                                CircuitLoweringState &loweringState);
   LogicalResult lowerFormalBody(verif::FormalOp formalOp,
                                 CircuitLoweringState &loweringState);
   LogicalResult lowerSimulationBody(verif::SimulationOp simulationOp,
                                     CircuitLoweringState &loweringState);
+  LogicalResult lowerBody(Operation *op, CircuitLoweringState &loweringState);
 };
 
 } // end anonymous namespace
@@ -617,9 +618,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                              verificationFlavor, getAnalysis<InstanceGraph>(),
                              &getAnalysis<NLATable>());
 
-  SmallVector<hw::HWModuleOp, 32> modulesToProcess;
-  SmallVector<verif::FormalOp> formalOpsToProcess;
-  SmallVector<verif::SimulationOp> simulationOpsToProcess;
+  SmallVector<Operation *, 32> opsToProcess;
 
   AnnotationSet circuitAnno(circuit);
   moveVerifAnno(getOperation(), circuitAnno, extractAssertAnnoClass,
@@ -643,7 +642,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                 return failure();
 
               state.recordModuleMapping(&op, loweredMod);
-              modulesToProcess.push_back(loweredMod);
+              opsToProcess.push_back(loweredMod);
               // Lower all the alias types.
               module.walk([&](Operation *op) {
                 for (auto res : op->getResults()) {
@@ -677,7 +676,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                   oldOp.getParametersAttr());
               newOp.getBody().emplaceBlock();
               state.recordModuleMapping(oldOp, newOp);
-              formalOpsToProcess.push_back(newOp);
+              opsToProcess.push_back(newOp);
               return success();
             })
             .Case<SimulationOp>([&](auto oldOp) {
@@ -689,7 +688,7 @@ void FIRRTLModuleLowering::runOnOperation() {
               body.addArgument(seq::ClockType::get(builder.getContext()), loc);
               body.addArgument(builder.getI1Type(), loc);
               state.recordModuleMapping(oldOp, newOp);
-              simulationOpsToProcess.push_back(newOp);
+              opsToProcess.push_back(newOp);
               return success();
             })
             .Default([&](Operation *op) {
@@ -747,28 +746,14 @@ void FIRRTLModuleLowering::runOnOperation() {
         ->setAttr(moduleHierarchyFileAttrName,
                   ArrayAttr::get(&getContext(), testHarnessHierarchyFiles));
 
-  // Lower all module bodies.
-  auto result = mlir::failableParallelForEachN(
-      &getContext(), 0, modulesToProcess.size(), [&](auto index) {
-        return lowerModuleOperations(modulesToProcess[index], state);
+  // Lower all module and formal op bodies.
+  auto result =
+      mlir::failableParallelForEach(&getContext(), opsToProcess, [&](auto op) {
+        return lowerBody(op, state);
       });
   if (failed(result))
     return signalPassFailure();
 
-  // Lower all formal op bodies.
-  result = mlir::failableParallelForEach(
-      &getContext(), formalOpsToProcess,
-      [&](auto op) { return lowerFormalBody(op, state); });
-  if (failed(result))
-    return signalPassFailure();
-
-  // Lower all simulation op bodies.
-  result = mlir::failableParallelForEach(
-      &getContext(), simulationOpsToProcess,
-      [&](auto op) { return lowerSimulationBody(op, state); });
-  if (failed(result))
-    return signalPassFailure();
-
   // Move binds from inside modules to outside modules.
   for (auto bind : state.binds) {
     bind->moveBefore(bind->getParentOfType<hw::HWModuleOp>());
@@ -1900,11 +1885,24 @@ struct FIRRTLLowering : public FIRRTLVisitor<FIRRTLLowering, LogicalResult> {
 };
 } // end anonymous namespace
 
-LogicalResult FIRRTLModuleLowering::lowerModuleOperations(
-    hw::HWModuleOp module, CircuitLoweringState &loweringState) {
+LogicalResult
+FIRRTLModuleLowering::lowerModuleBody(hw::HWModuleOp module,
+                                      CircuitLoweringState &loweringState) {
   return FIRRTLLowering(module, loweringState).run();
 }
 
+LogicalResult
+FIRRTLModuleLowering::lowerBody(Operation *op,
+                                CircuitLoweringState &loweringState) {
+  if (auto moduleOp = dyn_cast<hw::HWModuleOp>(op))
+    return lowerModuleBody(moduleOp, loweringState);
+  if (auto formalOp = dyn_cast<verif::FormalOp>(op))
+    return lowerFormalBody(formalOp, loweringState);
+  if (auto simulationOp = dyn_cast<verif::SimulationOp>(op))
+    return lowerSimulationBody(simulationOp, loweringState);
+  return failure();
+}
+
 // This is the main entrypoint for the lowering pass.
 LogicalResult FIRRTLLowering::run() {
   // Mark the module's block arguments are already lowered. This will allow