Parameters Performance Improvements (#177)

This PR adds performance improvements for working with parameters. 

Changes to MergeCircuitPass:
* Remove some unused vectors
* Change circuit merge so that only unique operands / arguments are
added when merging circuits rather than the union of the two circuits
arguments.

Author: bcdonovan

include/Dialect/QCS/IR/QCSOps.h

@@ -30,6 +30,7 @@
 #include "Dialect/QCS/IR/QCSTypes.h"
 
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/StringMap.h"
 
 #define GET_OP_CLASSES
 #include "Dialect/QCS/IR/QCSOps.h.inc"

include/Dialect/QCS/IR/QCSOps.td

@@ -289,7 +289,7 @@ def QCS_ParameterLoadOp : QCS_Op<"parameter_load",
 
     let extraClassDeclaration = [{
         // Return the initial value - using ParameterInitialValueAnalysis
-        ParameterType getInitialValue(std::unordered_map<std::string, ParameterType> &parameterNames);
+        ParameterType getInitialValue(llvm::StringMap<ParameterType> &parameterNames);
 
         // Return the initial value - slower SymbolTable version
         ParameterType getInitialValue();

include/Dialect/QCS/Utils/ParameterInitialValueAnalysis.h

@@ -20,10 +20,13 @@
 #define QCS_PARAMETER_INITIAL_VALUE_ANALYSIS_H
 
 #include "Dialect/QCS/IR/QCSOps.h"
+#include "HAL/SystemConfiguration.h"
 
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Pass/Pass.h"
 
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 
 #include <string>
@@ -33,16 +36,16 @@ namespace mlir::qcs {
 
 using namespace mlir;
 
+using InitialValueType = llvm::StringMap<ParameterType>;
+
 class ParameterInitialValueAnalysis {
 private:
-  std::unordered_map<std::string, ParameterType> initial_values_;
-  bool invalid_;
+  InitialValueType initial_values_;
+  bool invalid_{true};
 
 public:
   ParameterInitialValueAnalysis(mlir::Operation *op);
-  std::unordered_map<std::string, ParameterType> &getNames() {
-    return initial_values_;
-  }
+  InitialValueType &getNames() { return initial_values_; }
   void invalidate() { invalid_ = true; }
   bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
     return invalid_;

lib/Dialect/QCS/IR/QCSOps.cpp

@@ -29,6 +29,8 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
 
 using namespace mlir;
 using namespace mlir::qcs;
@@ -140,15 +142,15 @@ ParameterType ParameterLoadOp::getInitialValue() {
 }
 
 // Returns the float value from the initial value of this parameter
-// this version uses a precomputed map of parrameter_name to the intial_value
+// this version uses a precomputed map of parameter_name to the initial_value
 // in order to avoid slow SymbolTable lookups
 ParameterType ParameterLoadOp::getInitialValue(
-    std::unordered_map<std::string, ParameterType> &declareParametersMap) {
+    llvm::StringMap<ParameterType> &declareParametersMap) {
   auto *op = getOperation();
   auto paramRefAttr =
       op->getAttrOfType<mlir::FlatSymbolRefAttr>("parameter_name");
 
-  auto paramOpEntry = declareParametersMap.find(paramRefAttr.getValue().str());
+  auto paramOpEntry = declareParametersMap.find(paramRefAttr.getValue());
 
   if (paramOpEntry == declareParametersMap.end()) {
     op->emitError("Could not find declare parameter op" +

lib/Dialect/QCS/Utils/ParameterInitialValueAnalysis.cpp

@@ -27,27 +27,43 @@
 using namespace mlir::qcs;
 
 ParameterInitialValueAnalysis::ParameterInitialValueAnalysis(
-    mlir::Operation *op) {
-  op->walk([&](DeclareParameterOp declareParameterOp) {
-    double initial_value = 0.0;
-    if (declareParameterOp.initial_value().hasValue()) {
-      auto angleAttr = declareParameterOp.initial_value()
-                           .getValue()
-                           .dyn_cast<mlir::quir::AngleAttr>();
-      auto floatAttr =
-          declareParameterOp.initial_value().getValue().dyn_cast<FloatAttr>();
-      if (!(angleAttr || floatAttr))
-        op->emitError(
-            "Parameters are currently limited to angles or float[64] only.");
+    mlir::Operation *moduleOp) {
 
-      if (angleAttr)
-        initial_value = angleAttr.getValue().convertToDouble();
+  if (not invalid_)
+    return;
 
-      if (floatAttr)
-        initial_value = floatAttr.getValue().convertToDouble();
-    }
-    initial_values_[declareParameterOp.sym_name().str()] = initial_value;
-  });
+  // process the module top level to cache declareParameterOp initial_values
+  // this does not use a walk method so that submodule (if present) are not
+  // processed in order to limit processing time
+
+  for (auto &region : moduleOp->getRegions())
+    for (auto &block : region.getBlocks())
+      for (auto &op : block.getOperations()) {
+        auto declareParameterOp = dyn_cast<DeclareParameterOp>(op);
+        if (!declareParameterOp)
+          continue;
+
+        // moduleOp->walk([&](DeclareParameterOp declareParameterOp) {
+        double initial_value = 0.0;
+        if (declareParameterOp.initial_value().hasValue()) {
+          auto angleAttr = declareParameterOp.initial_value()
+                               .getValue()
+                               .dyn_cast<mlir::quir::AngleAttr>();
+          auto floatAttr = declareParameterOp.initial_value()
+                               .getValue()
+                               .dyn_cast<FloatAttr>();
+          if (!(angleAttr || floatAttr))
+            declareParameterOp.emitError("Parameters are currently limited to "
+                                         "angles or float[64] only.");
+
+          if (angleAttr)
+            initial_value = angleAttr.getValue().convertToDouble();
+
+          if (floatAttr)
+            initial_value = floatAttr.getValue().convertToDouble();
+        }
+        initial_values_[declareParameterOp.sym_name()] = initial_value;
+      }
   invalid_ = false;
 }
 
@@ -60,7 +76,7 @@ llvm::StringRef ParameterInitialValueAnalysisPass::getArgument() const {
 }
 
 llvm::StringRef ParameterInitialValueAnalysisPass::getDescription() const {
-  return "Run ParameterIntialValueAnalysis";
+  return "Run ParameterInitialValueAnalysis";
 }
 
 // TODO: move registerQCSPasses to separate source file if additional passes

lib/Dialect/QUIR/Transforms/MergeCircuits.cpp

@@ -33,6 +33,7 @@
 #include "llvm/ADT/StringRef.h"
 
 #include <algorithm>
+#include <unordered_map>
 #include <vector>
 
 using namespace mlir;
@@ -238,16 +239,29 @@ MergeCircuitsPass::mergeCallCircuits(PatternRewriter &rewriter,
 
   rewriter.setInsertionPointAfter(nextCircuitOp);
 
-  llvm::SmallVector<Type> inputTypes;
-  llvm::SmallVector<Value> inputValues;
   llvm::SmallVector<Type> outputTypes;
   llvm::SmallVector<Value> outputValues;
 
-  // merge input type into single SmallVector
-  inputTypes.append(circuitOp->getOperandTypes().begin(),
-                    circuitOp->getOperandTypes().end());
-  inputTypes.append(nextCircuitOp->getOperandTypes().begin(),
-                    nextCircuitOp->getOperandTypes().end());
+  // merge the call_circuits
+  // collect their input values
+  llvm::SmallVector<Value> callInputValues;
+  callInputValues.append(callCircuitOp->getOperands().begin(),
+                         callCircuitOp.getOperands().end());
+
+  llvm::SmallVector<int> insertedArguments;
+  std::unordered_map<int, int> reusedArguments;
+  int index = 0;
+  for (auto inputValue : nextCallCircuitOp->getOperands()) {
+    auto *search = find(callInputValues, inputValue);
+    if (search == callInputValues.end()) {
+      callInputValues.push_back(inputValue);
+      insertedArguments.push_back(index);
+    } else {
+      int originalIndex = search - callInputValues.begin();
+      reusedArguments[index] = originalIndex;
+    }
+    index++;
+  }
 
   // merge circuit names
   std::string newName =
@@ -269,12 +283,20 @@ MergeCircuitsPass::mergeCallCircuits(PatternRewriter &rewriter,
   // argument numbers
   BlockAndValueMapping mapper;
   auto baseArgNum = newCircuitOp.getNumArguments();
+  int insertedCount = 0;
   for (uint cnt = 0; cnt < nextCircuitOp.getNumArguments(); cnt++) {
     auto arg = nextCircuitOp.getArgument(cnt);
-    auto dictArg = nextCircuitOp.getArgAttrDict(cnt);
-    newCircuitOp.insertArgument(baseArgNum + cnt, arg.getType(), dictArg,
-                                arg.getLoc());
-    mapper.map(arg, newCircuitOp.getArgument(baseArgNum + cnt));
+    int argumentIndex = 0;
+    if (find(insertedArguments, cnt) != insertedArguments.end()) {
+      auto dictArg = nextCircuitOp.getArgAttrDict(cnt);
+      newCircuitOp.insertArgument(baseArgNum + insertedCount, arg.getType(),
+                                  dictArg, arg.getLoc());
+      argumentIndex = baseArgNum + insertedCount;
+      insertedCount++;
+    } else {
+      argumentIndex = reusedArguments[cnt];
+    }
+    mapper.map(arg, newCircuitOp.getArgument(argumentIndex));
   }
 
   // find return op in new circuit and copy second circuit into the
@@ -331,14 +353,6 @@ MergeCircuitsPass::mergeCallCircuits(PatternRewriter &rewriter,
   newCircuitOp->setAttr(mlir::quir::getPhysicalIdsAttrName(),
                         rewriter.getI32ArrayAttr(ArrayRef<int>(allIds)));
 
-  // merge the call_circuits
-  // collect their input values
-  llvm::SmallVector<Value> callInputValues;
-  callInputValues.append(callCircuitOp->getOperands().begin(),
-                         callCircuitOp.getOperands().end());
-  callInputValues.append(nextCallCircuitOp->getOperands().begin(),
-                         nextCallCircuitOp.getOperands().end());
-
   rewriter.setInsertionPointAfter(nextCallCircuitOp);
   auto newCallOp = rewriter.create<mlir::quir::CallCircuitOp>(
       callCircuitOp->getLoc(), newName, TypeRange(outputTypes),

test/Dialect/QUIR/Transforms/merge-circuits.mlir

@@ -110,21 +110,21 @@ module {
     %12:2 = quir.call_circuit @circuit_6(%0) : (!quir.qubit<1>) -> (i1, i1)
     quir.barrier %200 : (!quir.qubit<1>) -> ()
     %13:2 = quir.call_circuit @circuit_6(%0) : (!quir.qubit<1>) -> (i1, i1)
-    // CHECK: %{{.*}}:4 = quir.call_circuit @circuit_6_q0_circuit_6_q0(%0, %0) : (!quir.qubit<1>, !quir.qubit<1>) -> (i1, i1, i1, i1)
+    // CHECK: %{{.*}}:4 = quir.call_circuit @circuit_6_q0_circuit_6_q0(%0) : (!quir.qubit<1>) -> (i1, i1, i1, i1)
 
 
     quir.barrier %0, %1, %200, %201, %202 : (!quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>) -> ()
     %14:2 = quir.call_circuit @circuit_7(%0) : (!quir.qubit<1>) -> (i1, i1)
     quir.barrier %200, %201 : (!quir.qubit<1>, !quir.qubit<1>) -> ()
     quir.barrier %200, %202 : (!quir.qubit<1>, !quir.qubit<1>) -> ()
     %15:2 = quir.call_circuit @circuit_7(%0) : (!quir.qubit<1>) -> (i1, i1)
-    // CHECK: %{{.*}}:4 = quir.call_circuit @circuit_7_q0_circuit_7_q0(%0, %0) : (!quir.qubit<1>, !quir.qubit<1>) -> (i1, i1, i1, i1)
+    // CHECK: %{{.*}}:4 = quir.call_circuit @circuit_7_q0_circuit_7_q0(%0) : (!quir.qubit<1>) -> (i1, i1, i1, i1)
 
     quir.barrier %0, %1, %200, %201, %202 : (!quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>) -> ()
     %16:2 = quir.call_circuit @circuit_8(%0) : (!quir.qubit<1>) -> (i1, i1)
     quir.barrier %0 : (!quir.qubit<1>) -> ()
     %17:2 = quir.call_circuit @circuit_8(%0) : (!quir.qubit<1>) -> (i1, i1)
-    // CHECK-NOT: %{{.*}}:4 = quir.call_circuit @circuit_8_q0_circuit_8_q0(%0, %0) : (!quir.qubit<1>, !quir.qubit<1>) -> (i1, i1, i1, i1)
+    // CHECK-NOT: %{{.*}}:4 = quir.call_circuit @circuit_8_q0_circuit_8_q0(%0) : (!quir.qubit<1>) -> (i1, i1, i1, i1)
 
     %c0_i32 = arith.constant 0 : i32
     return %c0_i32 : i32