Merge branch 'main' into pr-merge-main-into-qubit-mgmt

.github/workflows/integration_tests.yml

@@ -617,7 +617,7 @@ jobs:
 
           # Test NVQC Python examples + Python MLIR execution tests (not IR tests)
           python3 -m pip install pytest
-          for ex in `find examples/python python/tests/mlir/target -name '*.py'`; do
+          for ex in `find examples/python python/tests/mlir/target -name '*.py' -not -path '*/python/tutorials/*'`; do
             filename=$(basename -- "$ex")
             filename="${filename%.*}"
             echo "Testing $filename:"
@@ -734,7 +734,7 @@ jobs:
           set +e # Allow script to keep going through errors
           python$python_version -m pip install pytest
           test_err_sum=0
-          for ex in `find examples/python python/tests/mlir/target -name '*.py'`; do
+          for ex in `find examples/python python/tests/mlir/target -name '*.py' -not -path '*/python/tutorials/*'`; do
             filename=$(basename -- "$ex")
             filename="${filename%.*}"
             echo "Testing $filename:"

.github/workflows/nvqc_regression_tests.yml

@@ -192,7 +192,7 @@ jobs:
 
           # Test NVQC Python examples + Python MLIR execution tests (not IR tests)
           python3 -m pip install pytest
-          for ex in `find examples/python python/tests/mlir/target -name '*.py'`; do
+          for ex in `find examples/python python/tests/mlir/target -name '*.py' -not -path '*/python/tutorials/*'`; do
             filename=$(basename -- "$ex")
             filename="${filename%.*}"
             echo "Testing $filename:"
@@ -306,7 +306,7 @@ jobs:
           set +e # Allow script to keep going through errors
           python$python_version -m pip install pytest
           test_err_sum=0
-          for ex in `find examples/python python/tests/mlir/target -name '*.py'`; do
+          for ex in `find examples/python python/tests/mlir/target -name '*.py' -not -path '*/python/tutorials/*'`; do
             filename=$(basename -- "$ex")
             filename="${filename%.*}"
             echo "Testing $filename:"

include/cudaq/Optimizer/Transforms/Passes.td

@@ -12,26 +12,27 @@
 include "mlir/Pass/PassBase.td"
 
 def AddWireset : Pass<"add-wireset", "mlir::ModuleOp"> {
-  let summary = "Adds a topology-less `quake.wireset` to the module";
+  let summary = "Adds a topology-less `quake.wire_set` to the module.";
   let description = [{
-    Adds a `quake.wireset` operation without tological info to the module.
+    Adds a `quake.wire_set` operation without any topological information to the
+    ModuleOp.
   }];
 }
 
 def AssignWireIndices : Pass<"assign-wire-indices", "mlir::func::FuncOp"> {
-  let summary = "Replaces wires with wires from a `quake.wireset`";
+  let summary = "Replaces `null_wire` with `borrow_wire` from a `wire_set`.";
   let description = [{
-    Replaces all instances of `quake.null_wire_op` with `quake.borrow_wire_op`s
-    from a common `quake.wireset` without any topological information.
-    Each wire is assigned a unique identifier (the index into the
-    `quake.wireset`) through this process.
+    Replaces all instances of `quake.null_wire` with `quake.borrow_wire` from a
+    common `quake.wire_set` without any topological information.  Each wire is
+    assigned a unique identifier (the index into the `quake.wire_set`) through
+    this process.
   }];
 }
 
 def ApplyControlNegations :
     Pass<"apply-control-negations", "mlir::func::FuncOp"> {
   let summary =
-    "Replace all operations with negative controls with positive controls and X operations.";
+    "Replace all `Ops` with negated controls with normal controls and `XOp`s.";
 
   let description = [{
     For every quantum operation with a negative control, replace that operation 

lib/Optimizer/Transforms/FactorQuantumAlloc.cpp

@@ -7,6 +7,7 @@
  ******************************************************************************/
 
 #include "PassDetails.h"
+#include "cudaq/Optimizer/Builder/Factory.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
 #include "cudaq/Optimizer/Transforms/Passes.h"
 #include "mlir/Transforms/DialectConversion.h"
@@ -42,21 +43,47 @@ class AllocaPat : public OpRewritePattern<quake::AllocaOp> {
     for (std::size_t i = 0; i < size; ++i)
       newAllocs.emplace_back(rewriter.create<quake::AllocaOp>(loc, refTy));
 
-    // 2. Visit all users and replace them accordingly.
-    for (auto *user : allocOp->getUsers()) {
-      if (auto dealloc = dyn_cast<quake::DeallocOp>(user)) {
-        rewriter.setInsertionPoint(dealloc);
-        auto deloc = dealloc.getLoc();
-        for (std::size_t i = 0; i < size - 1; ++i)
-          rewriter.create<quake::DeallocOp>(deloc, newAllocs[i]);
-        rewriter.replaceOpWithNewOp<quake::DeallocOp>(dealloc,
-                                                      newAllocs[size - 1]);
-        continue;
+    std::function<LogicalResult(Operation *, std::int64_t)> rewriteOpAndUsers =
+        [&](Operation *op, std::int64_t start) -> LogicalResult {
+      // First handle the users. Note that this can recurse.
+      for (auto *user : op->getUsers()) {
+        if (auto dealloc = dyn_cast<quake::DeallocOp>(user)) {
+          rewriter.setInsertionPoint(dealloc);
+          auto deloc = dealloc.getLoc();
+          for (std::size_t i = 0; i < size - 1; ++i)
+            rewriter.create<quake::DeallocOp>(deloc, newAllocs[i]);
+          rewriter.replaceOpWithNewOp<quake::DeallocOp>(dealloc,
+                                                        newAllocs[size - 1]);
+          continue;
+        }
+        if (auto subveq = dyn_cast<quake::SubVeqOp>(user)) {
+          auto lowInt = cudaq::opt::factory::getIntIfConstant(subveq.getLow());
+          if (!lowInt)
+            return failure();
+          for (auto *subUser : subveq->getUsers())
+            if (failed(rewriteOpAndUsers(subUser, *lowInt)))
+              return failure();
+          rewriter.eraseOp(subveq);
+          continue;
+        }
+        if (auto ext = dyn_cast<quake::ExtractRefOp>(user)) {
+          auto index = ext.getConstantIndex();
+          rewriter.replaceOp(ext, newAllocs[start + index].getResult());
+        }
       }
-      auto ext = cast<quake::ExtractRefOp>(user);
-      auto index = ext.getConstantIndex();
-      rewriter.replaceOp(ext, newAllocs[index].getResult());
-    }
+      // Now handle the base operation.
+      if (isa<quake::SubVeqOp>(op))
+        rewriter.eraseOp(op);
+      else if (auto ext = dyn_cast<quake::ExtractRefOp>(op)) {
+        auto index = ext.getConstantIndex();
+        rewriter.replaceOp(ext, newAllocs[start + index].getResult());
+      }
+      return success();
+    };
+
+    // 2. Visit all users and replace them accordingly.
+    if (failed(rewriteOpAndUsers(allocOp, 0)))
+      return failure();
 
     // 3. Remove the original alloca operation.
     rewriter.eraseOp(allocOp);
@@ -165,17 +192,32 @@ class FactorQuantumAllocationsPass
 
   LogicalResult runAnalysis(SmallVector<quake::AllocaOp> &allocations) {
     auto func = getOperation();
+    std::function<bool(Operation *)> isUseConvertible =
+        [&](Operation *op) -> bool {
+      if (isa<quake::DeallocOp>(op))
+        return true;
+      if (auto ext = dyn_cast<quake::ExtractRefOp>(op))
+        if (ext.hasConstantIndex())
+          return true;
+      if (auto sub = dyn_cast<quake::SubVeqOp>(op)) {
+        if (!cudaq::opt::factory::getIntIfConstant(sub.getLow()) ||
+            !cudaq::opt::factory::getIntIfConstant(sub.getHigh()))
+          return false;
+        for (auto *subUser : sub->getUsers())
+          if (!isUseConvertible(subUser))
+            return false;
+        return true;
+      }
+      return false;
+    };
     func.walk([&](quake::AllocaOp alloc) {
       if (!allocaOfVeq(alloc) || allocaOfUnspecifiedSize(alloc) ||
           alloc.hasInitializedState())
         return;
       bool usesAreConvertible = [&]() {
         for (auto *users : alloc->getUsers()) {
-          if (isa<quake::DeallocOp>(users))
+          if (isUseConvertible(users))
             continue;
-          if (auto ext = dyn_cast<quake::ExtractRefOp>(users))
-            if (ext.hasConstantIndex())
-              continue;
           return false;
         }
         return true;

lib/Optimizer/Transforms/LiftArrayAlloc.cpp

@@ -147,7 +147,7 @@ class AllocaPattern : public OpRewritePattern<cudaq::cc::AllocaOp> {
           rewriter.create<cudaq::cc::ConstantArrayOp>(loc, arrTy, valuesAttr);
     }
 
-    std::vector<mlir::Operation *> toErase;
+    SmallVector<Operation *> toErase;
 
     // Rewalk all the uses of alloc, u, which must be cc.cast or cc.compute_ptr.
     // For each,u, remove a store and replace a load with a cc.extract_value.

python/cudaq/runtime/sample.py

@@ -66,11 +66,10 @@ def sample(kernel, *args, shots_count=1000, noise_model=None):
     ctx = cudaq_runtime.ExecutionContext("sample", shots_count)
     cudaq_runtime.setExecutionContext(ctx)
 
-    if hasattr(kernel, 'metadata') and kernel.metadata['conditionalOnMeasure']:
+    if hasattr(kernel, 'metadata') and kernel.metadata.get(
+            'conditionalOnMeasure', False):
         ctx.hasConditionalsOnMeasureResults = True
 
-    cudaq_runtime.setExecutionContext(ctx)
-
     platformSupportsConditionalFeedback = cudaq_runtime.supportsConditionalFeedback(
     )
     if ctx.hasConditionalsOnMeasureResults and not platformSupportsConditionalFeedback:

runtime/common/ArgumentConversion.cpp

@@ -278,8 +278,10 @@ void cudaq::opt::ArgumentConverter::gen(const std::vector<void *> &arguments) {
   FunctionType fromFuncTy = fun.getFunctionType();
   for (auto iter :
        llvm::enumerate(llvm::zip(fromFuncTy.getInputs(), arguments))) {
-    Type argTy = std::get<0>(iter.value());
     void *argPtr = std::get<1>(iter.value());
+    if (!argPtr)
+      continue;
+    Type argTy = std::get<0>(iter.value());
     unsigned i = iter.index();
     auto buildSubst = [&, i = i]<typename... Ts>(Ts &&...ts) {
       builder.setInsertionPointToEnd(substModule.getBody());
@@ -360,3 +362,33 @@ void cudaq::opt::ArgumentConverter::gen(const std::vector<void *> &arguments) {
       substitutions.emplace_back(std::move(subst));
   }
 }
+
+void cudaq::opt::ArgumentConverter::gen(
+    const std::vector<void *> &arguments,
+    const std::unordered_set<unsigned> &exclusions) {
+  std::vector<void *> partialArgs;
+  for (auto iter : llvm::enumerate(arguments)) {
+    if (exclusions.contains(iter.index())) {
+      partialArgs.push_back(nullptr);
+      continue;
+    }
+    partialArgs.push_back(iter.value());
+  }
+  gen(partialArgs);
+}
+
+void cudaq::opt::ArgumentConverter::gen_drop_front(
+    const std::vector<void *> &arguments, unsigned numDrop) {
+  // If we're dropping all the arguments, we're done.
+  if (numDrop >= arguments.size())
+    return;
+  std::vector<void *> partialArgs;
+  for (void *arg : arguments) {
+    if (numDrop--) {
+      partialArgs.push_back(nullptr);
+      continue;
+    }
+    partialArgs.push_back(arg);
+  }
+  gen(partialArgs);
+}

runtime/common/ArgumentConversion.h

@@ -12,6 +12,7 @@
 #include "cudaq/Optimizer/Dialect/CC/CCTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Types.h"
+#include <unordered_set>
 
 namespace cudaq {
 class state;
@@ -29,6 +30,15 @@ class ArgumentConverter {
   /// The arguments are those presented to the kernel, kernelName.
   void gen(const std::vector<void *> &arguments);
 
+  /// Generate a substitution ModuleOp but include only the arguments that do
+  /// not appear in the set of \p exclusions.
+  void gen(const std::vector<void *> &arguments,
+           const std::unordered_set<unsigned> &exclusions);
+
+  /// Generate a substitution ModuleOp but drop the first \p numDrop arguments
+  /// and thereby exclude them from the substitutions.
+  void gen_drop_front(const std::vector<void *> &arguments, unsigned numDrop);
+
   /// Get the list of substitutions that were generated by `gen()`.
   mlir::SmallVector<cudaq::cc::ArgumentSubstitutionOp> &getSubstitutions() {
     return substitutions;

test/Quake/combine.qke

@@ -157,6 +157,15 @@ func.func @c() {
 // CHECK:           return
 // CHECK:         }
 
+// FACTOR-LABEL:  func.func @c() {
+// FACTOR:          %0 = quake.alloca !quake.ref
+// FACTOR:          %1 = quake.alloca !quake.ref
+// FACTOR:          %2 = quake.alloca !quake.ref
+// FACTOR:          %3 = quake.alloca !quake.ref
+// FACTOR:          quake.x %3 : (!quake.ref) -> ()
+// FACTOR:          return
+// FACTOR:        }
+
 func.func @d(%c2: i64, %c3: i64, %c1: i64, %off: i16) {
   %1 = quake.alloca !quake.veq<4>
   %2 = quake.subveq %1, %c2, %c3 : (!quake.veq<4>, i64, i64) -> !quake.veq<2>