Update parameter/constant op handling for performance (#335)

This PR updates how parameters are handled. They are now directly read
from "parameter_load" ops which removes the need for QUIR variable
analysis and casting which significantly improves performance of large
parameter programs by reducing the total number of operations in the
program. This also cuts down on the total number of arguments to each
circuit/sequence by storing the parameter loads/constants directly in
the sequence itself

For example before for the program below
```qasm
OPENQASM 3;

qubit $0;
qubit $1;
qubit $2;

gate sx q { }
gate rz(phi) q { }

input float[64] theta = 3.14159265358979;
input angle phi;

sx $0;
rz(theta) $0;

sx $0;

rz(phi) $1;
rz(3.141592) $1;

rz(theta) $2;
rz(phi) $2;


bit b;

b = measure $0;
```

The MLIR generated is now:
```mlir
module {
  func.func @sx(%arg0: !quir.qubit<1>) attributes {quir.classicalOnly = false} {
    return
  }
  func.func @rz(%arg0: !quir.qubit<1>, %arg1: !quir.angle<64>) attributes {quir.classicalOnly = false} {
    return
  }
  quir.circuit @circuit_0(%arg0: !quir.qubit<1> {quir.physicalId = 0 : i32}, %arg1: !quir.qubit<1> {quir.physicalId = 1 : i32}, %arg2: !quir.qubit<1> {quir.physicalId = 2 : i32}) -> i1 attributes {quir.classicalOnly = false, quir.physicalIds = [0 : i32, 1 : i32, 2 : i32]} {
    quir.call_gate @sx(%arg0) : (!quir.qubit<1>) -> ()
    %0 = qcs.parameter_load "theta" : !quir.angle<64> {initialValue = 3.14159265358979 : f64}
    quir.call_gate @rz(%arg0, %0) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    quir.call_gate @sx(%arg0) : (!quir.qubit<1>) -> ()
    %1 = qcs.parameter_load "phi" : !quir.angle<64>
    quir.call_gate @rz(%arg1, %1) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %angle = quir.constant #quir.angle<3.1415920000000002> : !quir.angle<64>
    quir.call_gate @rz(%arg1, %angle) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %2 = qcs.parameter_load "theta" : !quir.angle<64> {initialValue = 3.14159265358979 : f64}
    quir.call_gate @rz(%arg2, %2) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %3 = qcs.parameter_load "phi" : !quir.angle<64>
    quir.call_gate @rz(%arg2, %3) : (!quir.qubit<1>, !quir.angle<64>) -> ()
    %4 = quir.measure(%arg0) {quir.noFastPathComm, quir.noJunoComm, quir.noJunoUse} : (!quir.qubit<1>) -> i1
    quir.return %4 : i1
  }
  func.func @main() -> i32 attributes {quir.classicalOnly = false} {
    %c0_i32 = arith.constant 0 : i32
    %dur = quir.constant #quir.duration<4.500000e+06> : !quir.duration<dt>
    %c1 = arith.constant 1 : index
    %c1000 = arith.constant 1000 : index
    %c0 = arith.constant 0 : index
    qcs.init
    scf.for %arg0 = %c0 to %c1000 step %c1 {
      quir.delay %dur, () : !quir.duration<dt>, () -> ()
      qcs.shot_init {qcs.num_shots = 1000 : i32}
      %0 = quir.declare_qubit {id = 0 : i32} : !quir.qubit<1>
      %1 = quir.declare_qubit {id = 1 : i32} : !quir.qubit<1>
      %2 = quir.declare_qubit {id = 2 : i32} : !quir.qubit<1>
      %3 = quir.call_circuit @circuit_0(%0, %1, %2) : (!quir.qubit<1>, !quir.qubit<1>, !quir.qubit<1>) -> i1
    } {qcs.shot_loop, quir.classicalOnly = false, quir.physicalIds = [0 : i32, 1 : i32, 2 : i32]}
    qcs.finalize
    return %c0_i32 : i32
  }
}
```

Author: taalexander

conan/qasm/conandata.yml

@@ -1,5 +1,5 @@
 sources:
-  hash: "f6d695fd9f18462e65f6290d05ccb4ccb371b288"
+  hash: "ec7731bf645240a597cd9ebb2c395b114f155ed2"
 requirements:
   - "gmp/6.3.0"
   - "mpfr/4.1.0"

conan/qasm/conanfile.py

@@ -17,7 +17,7 @@
 
 class QasmConan(ConanFile):
     name = "qasm"
-    version = "0.3.2"
+    version = "0.3.3"
     url = "https://github.com/openqasm/qe-qasm.git"
     settings = "os", "compiler", "build_type", "arch"
     options = {"shared": [True, False], "examples": [True, False]}

conandata.yml

@@ -7,4 +7,4 @@ requirements:
   - pybind11/2.11.1
   - clang-tools-extra/17.0.5-0@
   - llvm/17.0.5-0@
-  - qasm/0.3.2@qss/stable
+  - qasm/0.3.3@qss/stable

include/Conversion/QUIRToPulse/QUIRToPulse.h

@@ -127,7 +127,7 @@ struct QUIRToPulsePass
                                    mlir::func::FuncOp &mainFunc);
   // map of the hashed location of quir angle/duration ops to their converted
   // pulse ops
-  std::unordered_map<std::string, mlir::Value>
+  std::unordered_map<Operation *, mlir::Value>
       classicalQUIROpLocToConvertedPulseOpMap;
 
   // port name to Port_CreateOp map

include/Dialect/QUIR/Transforms/ExtractCircuits.h

@@ -25,11 +25,11 @@
 #include "Utils/SymbolCacheAnalysis.h"
 
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
-#include "llvm/ADT/SmallVector.h"
-
+#include <set>
 #include <unordered_map>
 
 namespace mlir::quir {
@@ -49,14 +49,14 @@ struct ExtractCircuitsPass
                     OpBuilder circuitBuilder);
   OpBuilder startCircuit(mlir::Location location, OpBuilder topLevelBuilder);
   void endCircuit(mlir::Operation *firstOp, mlir::Operation *lastOp,
-                  OpBuilder topLevelBuilder, OpBuilder circuitBuilder,
-                  llvm::SmallVector<Operation *> &eraseList);
-  void addToCircuit(mlir::Operation *currentOp, OpBuilder circuitBuilder,
-                    llvm::SmallVector<Operation *> &eraseList);
+                  OpBuilder topLevelBuilder, OpBuilder circuitBuilder);
+  void addToCircuit(mlir::Operation *currentOp, OpBuilder circuitBuilder);
+
   uint64_t circuitCount = 0;
   qssc::utils::SymbolCacheAnalysis *symbolCache{nullptr};
 
   mlir::quir::CircuitOp currentCircuitOp = nullptr;
+  mlir::IRMapping currentCircuitMapper;
   mlir::quir::CallCircuitOp newCallCircuitOp;
 
   llvm::SmallVector<Type> inputTypes;
@@ -68,6 +68,8 @@ struct ExtractCircuitsPass
 
   std::unordered_map<Operation *, uint32_t> circuitOperands;
   llvm::SmallVector<OpResult> originalResults;
+  std::set<Operation *> eraseConstSet;
+  std::set<Operation *> eraseOpSet;
 
 }; // struct ExtractCircuitsPass
 } // namespace mlir::quir

include/Frontend/OpenQASM3/QUIRGenQASM3Visitor.h

@@ -321,6 +321,9 @@ class QUIRGenQASM3Visitor : public BaseQASM3Visitor {
   mlir::Type getQUIRTypeFromDeclaration(const QASM::ASTDeclarationNode *);
 
   bool enableParametersWarningEmitted = false;
+
+  /// Cached dummy value for error handling
+  mlir::Value voidValue;
 };
 
 } // namespace qssc::frontend::openqasm3

include/Frontend/OpenQASM3/QUIRVariableBuilder.h

@@ -68,7 +68,7 @@ class QUIRVariableBuilder {
 
   mlir::Value generateParameterLoad(mlir::Location location,
                                     llvm::StringRef variableName,
-                                    mlir::Value assignedValue);
+                                    double initialValue);
 
   /// Generate code for declaring an array (at the builder's current insertion
   /// point).

lib/Conversion/QUIRToPulse/LoadPulseCals.cpp

@@ -152,7 +152,7 @@ void LoadPulseCalsPass::loadPulseCals(CallCircuitOp callCircuitOp,
       LLVM_DEBUG(llvm::dbgs() << "no pulse cal loading needed for " << op);
       assert((!op->hasTrait<mlir::quir::UnitaryOp>() and
               !op->hasTrait<mlir::quir::CPTPOp>()) &&
-             "unkown operation");
+             "unknown operation");
     }
   });
 }

lib/Conversion/QUIRToPulse/QUIRToPulse.cpp

@@ -100,8 +100,10 @@ void QUIRToPulsePass::runOnOperation() {
   moduleOp->walk([&](CallCircuitOp callCircOp) {
     if (isa<CircuitOp>(callCircOp->getParentOp()))
       return;
+
     auto convertedPulseCallSequenceOp =
         convertCircuitToSequence(callCircOp, mainFunc, moduleOp);
+
     if (!callCircOp->use_empty())
       callCircOp->replaceAllUsesWith(convertedPulseCallSequenceOp);
     callCircOp->erase();
@@ -229,8 +231,9 @@ QUIRToPulsePass::convertCircuitToSequence(CallCircuitOp &callCircuitOp,
       auto *newDelayCyclesOp = builder.clone(*quirOp);
       newDelayCyclesOp->moveAfter(callCircuitOp);
     } else
-      assert(((isa<quir::ConstantOp>(quirOp) or isa<quir::ReturnOp>(quirOp) or
-               isa<quir::CircuitOp>(quirOp))) &&
+      assert(((isa<quir::ConstantOp>(quirOp) ||
+               isa<qcs::ParameterLoadOp>(quirOp) ||
+               isa<quir::ReturnOp>(quirOp) || isa<quir::CircuitOp>(quirOp))) &&
              "quir op is not allowed in this pass.");
   });
 
@@ -251,6 +254,7 @@ QUIRToPulsePass::convertCircuitToSequence(CallCircuitOp &callCircuitOp,
                                                   convertedPulseSequenceOp,
                                                   convertedPulseSequenceOpArgs);
   convertedPulseCallSequenceOp->moveAfter(callCircuitOp);
+
   return convertedPulseCallSequenceOp;
 }
 
@@ -286,7 +290,7 @@ void QUIRToPulsePass::processCircuitArgs(
     } else if (argumentType.isa<mlir::quir::QubitType>()) {
       auto *qubitOp = callCircuitOp.getOperand(cnt).getDefiningOp();
     } else
-      llvm_unreachable("unkown circuit argument.");
+      llvm_unreachable("unknown circuit argument.");
   }
 }
 
@@ -339,7 +343,7 @@ void QUIRToPulsePass::processPulseCalArgs(
     } else if (argumentType.isa<FloatType>()) {
       assert(argAttr[index].dyn_cast<StringAttr>().getValue().str() ==
                  "angle" &&
-             "unkown argument.");
+             "unknown argument.");
       assert(angleOperands.size() && "no angle operand found.");
       auto nextAngle = angleOperands.front();
       LLVM_DEBUG(llvm::dbgs() << "angle argument ");
@@ -350,7 +354,7 @@ void QUIRToPulsePass::processPulseCalArgs(
     } else if (argumentType.isa<IntegerType>()) {
       assert(argAttr[index].dyn_cast<StringAttr>().getValue().str() ==
                  "duration" &&
-             "unkown argument.");
+             "unknown argument.");
       assert(durationOperands.size() && "no duration operand found.");
       auto nextDuration = durationOperands.front();
       LLVM_DEBUG(llvm::dbgs() << "duration argument ");
@@ -359,7 +363,7 @@ void QUIRToPulsePass::processPulseCalArgs(
                          pulseCalSequenceArgs, builder);
       durationOperands.pop();
     } else
-      llvm_unreachable("unkown argument type.");
+      llvm_unreachable("unknown argument type.");
   }
 }
 
@@ -379,12 +383,13 @@ void QUIRToPulsePass::getQUIROpClassicalOperands(
   }
 
   for (auto operand : classicalOperands)
-    if (operand.getType().isa<mlir::quir::AngleType>())
+    if (operand.getType().isa<mlir::quir::AngleType>() ||
+        operand.getType().isa<FloatType>())
       angleOperands.push(operand);
     else if (operand.getType().isa<mlir::quir::DurationType>())
       durationOperands.push(operand);
     else
-      llvm_unreachable("unkown operand.");
+      llvm_unreachable("unknown operand.");
 }
 
 void QUIRToPulsePass::processMixFrameOpArg(
@@ -463,21 +468,38 @@ void QUIRToPulsePass::processAngleArg(Value nextAngleOperand,
     pulseCalSequenceArgs.push_back(
         convertedPulseSequenceOp
             .getArguments()[circuitArgToConvertedSequenceArgMap[circNum]]);
-  } else {
-    auto angleOp = nextAngleOperand.getDefiningOp<mlir::quir::ConstantOp>();
-    std::string const angleLocHash =
-        std::to_string(mlir::hash_value(angleOp->getLoc()));
-    if (classicalQUIROpLocToConvertedPulseOpMap.find(angleLocHash) ==
+  } else if (auto angleOp =
+                 nextAngleOperand.getDefiningOp<mlir::quir::ConstantOp>()) {
+    auto *op = angleOp.getOperation();
+    if (classicalQUIROpLocToConvertedPulseOpMap.find(op) ==
         classicalQUIROpLocToConvertedPulseOpMap.end()) {
       double const angleVal =
           angleOp.getAngleValueFromConstant().convertToDouble();
       auto f64Angle = entryBuilder.create<mlir::arith::ConstantOp>(
           angleOp.getLoc(), entryBuilder.getFloatAttr(entryBuilder.getF64Type(),
                                                       llvm::APFloat(angleVal)));
-      classicalQUIROpLocToConvertedPulseOpMap[angleLocHash] = f64Angle;
+      classicalQUIROpLocToConvertedPulseOpMap[op] = f64Angle;
     }
-    pulseCalSequenceArgs.push_back(
-        classicalQUIROpLocToConvertedPulseOpMap[angleLocHash]);
+    pulseCalSequenceArgs.push_back(classicalQUIROpLocToConvertedPulseOpMap[op]);
+  } else if (auto paramOp =
+                 nextAngleOperand.getDefiningOp<mlir::qcs::ParameterLoadOp>()) {
+    auto *op = paramOp.getOperation();
+    if (classicalQUIROpLocToConvertedPulseOpMap.find(op) ==
+        classicalQUIROpLocToConvertedPulseOpMap.end()) {
+
+      auto newParam = entryBuilder.create<qcs::ParameterLoadOp>(
+          paramOp->getLoc(), entryBuilder.getF64Type(),
+          paramOp.getParameterName());
+      if (paramOp->hasAttr("initialValue")) {
+        auto initAttr = paramOp->getAttr("initialValue").dyn_cast<FloatAttr>();
+        if (initAttr)
+          newParam->setAttr("initialValue", initAttr);
+      }
+
+      classicalQUIROpLocToConvertedPulseOpMap[op] = newParam;
+    }
+
+    pulseCalSequenceArgs.push_back(classicalQUIROpLocToConvertedPulseOpMap[op]);
   }
 }
 
@@ -501,25 +523,23 @@ void QUIRToPulsePass::processDurationArg(
                TimeUnits::dt &&
            "this pass only accepts durations with dt unit");
 
-    if (classicalQUIROpLocToConvertedPulseOpMap.find(durLocHash) ==
+    auto *op = durationOp.getOperation();
+    if (classicalQUIROpLocToConvertedPulseOpMap.find(op) ==
         classicalQUIROpLocToConvertedPulseOpMap.end()) {
       auto dur64 = entryBuilder.create<mlir::arith::ConstantOp>(
           durationOp.getLoc(),
           entryBuilder.getIntegerAttr(entryBuilder.getI64Type(),
                                       uint64_t(durVal)));
-      classicalQUIROpLocToConvertedPulseOpMap[durLocHash] = dur64;
+      classicalQUIROpLocToConvertedPulseOpMap[op] = dur64;
     }
-    pulseCalSequenceArgs.push_back(
-        classicalQUIROpLocToConvertedPulseOpMap[durLocHash]);
+    pulseCalSequenceArgs.push_back(classicalQUIROpLocToConvertedPulseOpMap[op]);
   }
 }
 
 mlir::Value QUIRToPulsePass::convertAngleToF64(Operation *angleOp,
                                                mlir::OpBuilder &builder) {
   assert(angleOp && "angle op is null");
-  std::string const angleLocHash =
-      std::to_string(mlir::hash_value(angleOp->getLoc()));
-  if (classicalQUIROpLocToConvertedPulseOpMap.find(angleLocHash) ==
+  if (classicalQUIROpLocToConvertedPulseOpMap.find(angleOp) ==
       classicalQUIROpLocToConvertedPulseOpMap.end()) {
     if (auto castOp = dyn_cast<quir::ConstantOp>(angleOp)) {
       double const angleVal =
@@ -528,41 +548,46 @@ mlir::Value QUIRToPulsePass::convertAngleToF64(Operation *angleOp,
           castOp->getLoc(),
           builder.getFloatAttr(builder.getF64Type(), llvm::APFloat(angleVal)));
       f64Angle->moveAfter(castOp);
-      classicalQUIROpLocToConvertedPulseOpMap[angleLocHash] = f64Angle;
+      classicalQUIROpLocToConvertedPulseOpMap[angleOp] = f64Angle;
     } else if (auto castOp = dyn_cast<qcs::ParameterLoadOp>(angleOp)) {
-      auto angleCastedOp = builder.create<oq3::CastOp>(
-          castOp->getLoc(), builder.getF64Type(), castOp.getRes());
-      angleCastedOp->moveAfter(castOp);
-      classicalQUIROpLocToConvertedPulseOpMap[angleLocHash] = angleCastedOp;
+      // Just convert to an f64 directly
+      auto newParam = builder.create<qcs::ParameterLoadOp>(
+          angleOp->getLoc(), builder.getF64Type(), castOp.getParameterName());
+      if (castOp->hasAttr("initialValue")) {
+        auto initAttr = castOp->getAttr("initialValue").dyn_cast<FloatAttr>();
+        if (initAttr)
+          newParam->setAttr("initialValue", initAttr);
+      }
+      newParam->moveAfter(castOp);
+
+      classicalQUIROpLocToConvertedPulseOpMap[angleOp] = newParam;
     } else if (auto castOp = dyn_cast<oq3::CastOp>(angleOp)) {
       auto castOpArg = castOp.getArg();
       if (auto paramCastOp =
               dyn_cast<qcs::ParameterLoadOp>(castOpArg.getDefiningOp())) {
         auto angleCastedOp = builder.create<oq3::CastOp>(
             paramCastOp->getLoc(), builder.getF64Type(), paramCastOp.getRes());
         angleCastedOp->moveAfter(paramCastOp);
-        classicalQUIROpLocToConvertedPulseOpMap[angleLocHash] = angleCastedOp;
+        classicalQUIROpLocToConvertedPulseOpMap[angleOp] = angleCastedOp;
       } else if (auto constOp =
                      dyn_cast<arith::ConstantOp>(castOpArg.getDefiningOp())) {
         // if cast from float64 then use directly
         assert(constOp.getType() == builder.getF64Type() &&
                "expected angle type to be float 64");
-        classicalQUIROpLocToConvertedPulseOpMap[angleLocHash] = constOp;
+        classicalQUIROpLocToConvertedPulseOpMap[angleOp] = constOp;
       } else
         llvm_unreachable("castOp arg unknown");
     } else
       llvm_unreachable("angleOp unknown");
   }
-  return classicalQUIROpLocToConvertedPulseOpMap[angleLocHash];
+  return classicalQUIROpLocToConvertedPulseOpMap[angleOp];
 }
 
 mlir::Value QUIRToPulsePass::convertDurationToI64(
     mlir::quir::CallCircuitOp &callCircuitOp, Operation *durationOp, uint &cnt,
     mlir::OpBuilder &builder, mlir::func::FuncOp &mainFunc) {
   assert(durationOp && "duration op is null");
-  std::string const durLocHash =
-      std::to_string(mlir::hash_value(durationOp->getLoc()));
-  if (classicalQUIROpLocToConvertedPulseOpMap.find(durLocHash) ==
+  if (classicalQUIROpLocToConvertedPulseOpMap.find(durationOp) ==
       classicalQUIROpLocToConvertedPulseOpMap.end()) {
     if (auto castOp = dyn_cast<quir::ConstantOp>(durationOp)) {
       auto durVal =
@@ -575,11 +600,11 @@ mlir::Value QUIRToPulsePass::convertDurationToI64(
           castOp->getLoc(),
           builder.getIntegerAttr(builder.getI64Type(), uint64_t(durVal)));
       I64Dur->moveAfter(castOp);
-      classicalQUIROpLocToConvertedPulseOpMap[durLocHash] = I64Dur;
+      classicalQUIROpLocToConvertedPulseOpMap[durationOp] = I64Dur;
     } else
-      llvm_unreachable("unkown duration op");
+      llvm_unreachable("unknown duration op");
   }
-  return classicalQUIROpLocToConvertedPulseOpMap[durLocHash];
+  return classicalQUIROpLocToConvertedPulseOpMap[durationOp];
 }
 
 mlir::pulse::Port_CreateOp

lib/Dialect/Pulse/IR/PulseOps.cpp

@@ -17,6 +17,7 @@
 #include "Dialect/Pulse/IR/PulseOps.h"
 
 #include "Dialect/Pulse/IR/PulseTraits.h"
+#include "Dialect/QCS/IR/QCSOps.h"
 #include "Dialect/QUIR/IR/QUIROps.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -356,8 +357,9 @@ LogicalResult verifyClassical_(SequenceOp op) {
   mlir::Operation *classicalOp = nullptr;
   WalkResult const result = op->walk([&](Operation *subOp) {
     if (isa<mlir::arith::ConstantOp>(subOp) || isa<quir::ConstantOp>(subOp) ||
-        isa<CallSequenceOp>(subOp) || isa<pulse::ReturnOp>(subOp) ||
-        isa<SequenceOp>(subOp) || isa<mlir::complex::CreateOp>(subOp) ||
+        isa<qcs::ParameterLoadOp>(subOp) || isa<CallSequenceOp>(subOp) ||
+        isa<pulse::ReturnOp>(subOp) || isa<SequenceOp>(subOp) ||
+        isa<mlir::complex::CreateOp>(subOp) ||
         subOp->hasTrait<mlir::pulse::SequenceAllowed>() ||
         subOp->hasTrait<mlir::pulse::SequenceRequired>())
       return WalkResult::advance();

lib/Dialect/Pulse/Transforms/Scheduling.cpp

@@ -112,6 +112,7 @@ void QuantumCircuitPulseSchedulingPass::scheduleAlap(
             opEnd = quantumCircuitSequenceOpBlock->rend();
        opIt != opEnd; ++opIt) {
     auto &op = *opIt;
+
     if (auto quantumGateCallSequenceOp =
             dyn_cast<mlir::pulse::CallSequenceOp>(op)) {
       // find quantum gate SequenceOp

lib/Dialect/QUIR/IR/QUIROps.cpp

@@ -380,8 +380,9 @@ LogicalResult verifyClassical_(CircuitOp op) {
   mlir::Operation *classicalOp = nullptr;
   WalkResult const result = op->walk([&](Operation *subOp) {
     if (isa<mlir::arith::ConstantOp>(subOp) || isa<quir::ConstantOp>(subOp) ||
-        isa<CallCircuitOp>(subOp) || isa<quir::ReturnOp>(subOp) ||
-        isa<CircuitOp>(subOp) || subOp->hasTrait<mlir::quir::UnitaryOp>() ||
+        isa<qcs::ParameterLoadOp>(subOp) || isa<CallCircuitOp>(subOp) ||
+        isa<quir::ReturnOp>(subOp) || isa<CircuitOp>(subOp) ||
+        subOp->hasTrait<mlir::quir::UnitaryOp>() ||
         subOp->hasTrait<mlir::quir::CPTPOp>())
       return WalkResult::advance();
     classicalOp = subOp;