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 + } +}