Extract Circuits Performance Improvement

include/Dialect/QUIR/Transforms/ExtractCircuits.h

@@ -23,6 +23,7 @@
 
 #include "Dialect/QUIR/IR/QUIROps.h"
 
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
@@ -60,8 +61,8 @@ struct ExtractCircuitsPass
   llvm::SmallVector<Type> outputTypes;
   llvm::SmallVector<Value> outputValues;
   std::vector<int> phyiscalIds;
+  std::unordered_map<uint32_t, int> argToId;
 
-  std::unordered_map<uint32_t, BlockArgument> circuitArguments;
   std::unordered_map<Operation *, uint32_t> circuitOperands;
   llvm::SmallVector<OpResult> originalResults;
 

lib/Dialect/QUIR/Transforms/ExtractCircuits.cpp

@@ -91,9 +91,9 @@ OpBuilder ExtractCircuitsPass::startCircuit(Location location,
   outputTypes.clear();
   outputValues.clear();
   originalResults.clear();
-  circuitArguments.clear();
   circuitOperands.clear();
   phyiscalIds.clear();
+  argToId.clear();
 
   std::string const circuitName = "circuit_";
   std::string newName = circuitName + std::to_string(circuitCount++);
@@ -132,19 +132,14 @@ void ExtractCircuitsPass::addToCircuit(
     if (search == circuitOperands.end()) {
       argumentIndex = inputValues.size();
       inputValues.push_back(operand);
+      inputTypes.push_back(operand.getType());
       circuitOperands[defOp] = argumentIndex;
-
-      currentCircuitOp.insertArgument(argumentIndex, operand.getType(), {},
-                                      currentOp->getLoc());
+      currentCircuitOp.getBody().addArgument(operand.getType(),
+                                             currentOp->getLoc());
       if (isa<quir::DeclareQubitOp>(defOp)) {
-        auto physicalId = defOp->getAttrOfType<IntegerAttr>("id");
-        phyiscalIds.push_back(physicalId.getInt());
-        currentCircuitOp.setArgAttrs(
-            argumentIndex,
-            ArrayRef({NamedAttribute(
-                StringAttr::get(&getContext(),
-                                mlir::quir::getPhysicalIdAttrName()),
-                physicalId)}));
+        auto id = defOp->getAttrOfType<IntegerAttr>("id").getInt();
+        phyiscalIds.push_back(id);
+        argToId[argumentIndex] = id;
       }
     } else {
       argumentIndex = search->second;
@@ -175,17 +170,25 @@ void ExtractCircuitsPass::endCircuit(
   // change the input / output types for the quir.circuit
   auto opType = currentCircuitOp.getFunctionType();
   currentCircuitOp.setType(topLevelBuilder.getFunctionType(
-      /*inputs=*/opType.getInputs(),
+      /*inputs=*/ArrayRef<Type>(inputTypes),
       /*results=*/ArrayRef<Type>(outputTypes)));
 
+  for (const auto &[key, value] : argToId)
+    currentCircuitOp.setArgAttrs(
+        key,
+        ArrayRef({NamedAttribute(
+            StringAttr::get(&getContext(), mlir::quir::getPhysicalIdAttrName()),
+            topLevelBuilder.getI32IntegerAttr(value))}));
+
   std::sort(phyiscalIds.begin(), phyiscalIds.end());
   currentCircuitOp->setAttr(
       mlir::quir::getPhysicalIdsAttrName(),
       topLevelBuilder.getI32ArrayAttr(ArrayRef<int>(phyiscalIds)));
 
   // insert call_circuit
   // NOLINTNEXTLINE(misc-const-correctness)
-  OpBuilder builder(firstOp);
+  OpBuilder builder(lastOp);
+  builder.setInsertionPointAfter(lastOp);
   newCallCircuitOp = builder.create<mlir::quir::CallCircuitOp>(
       currentCircuitOp->getLoc(), currentCircuitOp.getSymName(),
       TypeRange(outputTypes), ValueRange(inputValues));