diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fb61df48bf..b6f6f5e0f2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -190,6 +190,10 @@
* Fixed the assembly format for `quantum.adjoint` when it has no quantum operands/results.
[(#2938)](https://github.com/PennyLaneAI/catalyst/pull/2938)
+* The `--decompose-lowering` pass no longer crashes on qreg-mode decomposition rules whose wire
+ indices are passed as separate scalar arguments, e.g. `(qreg, *params, idx0, idx1, ...)`.
+ [(#2952)](https://github.com/PennyLaneAI/catalyst/pull/2952)
+
Internal changes ⚙️
* The `/benchmark` GitHub comment trigger can now accept additional arguments and has been renamed to `!benchmark`.
diff --git a/mlir/lib/Quantum/Transforms/DecomposeLoweringImpl.hpp b/mlir/lib/Quantum/Transforms/DecomposeLoweringImpl.hpp
index 4fdc0d421c..60c9b7cd25 100644
--- a/mlir/lib/Quantum/Transforms/DecomposeLoweringImpl.hpp
+++ b/mlir/lib/Quantum/Transforms/DecomposeLoweringImpl.hpp
@@ -220,11 +220,26 @@ class BaseSignatureAnalyzer {
}
if (hasQreg) {
+ // The wire indices may be packed into a single tensor argument (e.g.
+ // tensor<2xi64>) or spread across several scalar arguments, as in
+ // signatures of the form (qreg, *params, first_qubit_index,
+ // second_qubit_index, ...). Consume the function inputs one at a time,
+ // letting the number of elements each input type expects decide how many
+ // wire indices it absorbs.
for (const auto &indices : {signature.inWireIndices, signature.inCtrlWireIndices}) {
- if (!indices.empty()) {
+ size_t consumed = 0;
+ while (consumed < indices.size()) {
+ assert(operandIdx < static_cast(funcInputsNoQreg.size()) &&
+ "ran out of function inputs while mapping wire indices");
+ Type funcInputType = funcInputsNoQreg[operandIdx];
+ size_t numToConsume =
+ std::min(getElementsCount(funcInputType), indices.size() - consumed);
+ ArrayRef indexSlice =
+ ArrayRef(indices).slice(consumed, numToConsume);
operands[operandIdx] =
- fromTensorOrAsIs(indices, funcInputsNoQreg[operandIdx], rewriter, loc);
+ fromTensorOrAsIs(indexSlice, funcInputType, rewriter, loc);
operandIdx++;
+ consumed += numToConsume;
}
}
}
diff --git a/mlir/lib/Quantum/Transforms/DecomposeLoweringPatterns.cpp b/mlir/lib/Quantum/Transforms/DecomposeLoweringPatterns.cpp
index ae8fd55c15..a4f0fdc30d 100644
--- a/mlir/lib/Quantum/Transforms/DecomposeLoweringPatterns.cpp
+++ b/mlir/lib/Quantum/Transforms/DecomposeLoweringPatterns.cpp
@@ -211,7 +211,6 @@ struct DLPauliRotOpPattern : public OpRewritePattern {
return failure();
}
func::FuncOp decompFunc = it->second;
- decompFunc.dump();
// Here is the assumption that the decomposition function must have at least one input
// and one result
@@ -228,9 +227,6 @@ struct DLPauliRotOpPattern : public OpRewritePattern {
assert(analyzer && "Analyzer should be valid");
auto callOperands = analyzer.prepareCallOperands(decompFunc, rewriter, op.getLoc());
- for (auto operand : callOperands) {
- operand.dump();
- }
auto callOp =
func::CallOp::create(rewriter, op.getLoc(), decompFunc.getFunctionType().getResults(),
decompFunc.getSymName(), callOperands);
diff --git a/mlir/test/Quantum/DecomposeLoweringTest.mlir b/mlir/test/Quantum/DecomposeLoweringTest.mlir
index b7abe6f203..94e73e38d0 100644
--- a/mlir/test/Quantum/DecomposeLoweringTest.mlir
+++ b/mlir/test/Quantum/DecomposeLoweringTest.mlir
@@ -792,3 +792,42 @@ module @different_qreg_values{
return %7 : !quantum.reg
}
}
+
+// -----
+
+// Test a qreg-mode decomposition rule whose wire indices are passed as separate
+// scalar arguments, i.e. a signature of the form
+// (qreg, *params, first_qubit_index, second_qubit_index, ...) instead of packing
+// the indices into a single tensor argument.
+module @paulirot_scalar_wire_indices {
+ func.func public @circuit(%arg0: tensor, %arg1: tensor) attributes {quantum.node} {
+ %cst = arith.constant 2.000000e-01 : f64
+ %0 = quantum.alloc( 2) : !quantum.reg
+ %extracted = tensor.extract %arg0[] : tensor
+ %1 = quantum.extract %0[%extracted] : !quantum.reg -> !quantum.bit
+ %extracted_0 = tensor.extract %arg1[] : tensor
+ %2 = quantum.extract %0[%extracted_0] : !quantum.reg -> !quantum.bit
+
+ // CHECK-NOT: quantum.paulirot
+ // CHECK: quantum.multirz
+ %out_qubits:2 = quantum.paulirot ["X", "Z"](%cst) %1, %2 : !quantum.bit, !quantum.bit
+
+ %3 = quantum.insert %0[%extracted], %out_qubits#0 : !quantum.reg, !quantum.bit
+ %4 = quantum.insert %3[%extracted_0], %out_qubits#1 : !quantum.reg, !quantum.bit
+ quantum.dealloc %4 : !quantum.reg
+ return
+ }
+
+ // CHECK-NOT: func.func private @my_paulirot_xz_decomp
+ func.func private @my_paulirot_xz_decomp(%inreg: !quantum.reg, %angle: tensor, %q0: tensor, %q1: tensor) -> !quantum.reg attributes {target_gate = "paulirotXZ"} {
+ %angle_f = tensor.extract %angle[] : tensor
+ %i0 = tensor.extract %q0[] : tensor
+ %i1 = tensor.extract %q1[] : tensor
+ %qb0 = quantum.extract %inreg[%i0] : !quantum.reg -> !quantum.bit
+ %qb1 = quantum.extract %inreg[%i1] : !quantum.reg -> !quantum.bit
+ %m:2 = quantum.multirz(%angle_f) %qb0, %qb1 : !quantum.bit, !quantum.bit
+ %r0 = quantum.insert %inreg[%i0], %m#0 : !quantum.reg, !quantum.bit
+ %r1 = quantum.insert %r0[%i1], %m#1 : !quantum.reg, !quantum.bit
+ return %r1 : !quantum.reg
+ }
+}