Skip to content

Canonicalize extract past distinct-index insert #2965

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rylanmalarchick wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Canonicalize extract past distinct-index insert #2965

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
449e1a7
Canonicalize extract past distinct-index insert
rylanmalarchick Jun 18, 2026
e9c448e
Reference PR #2965 in changelog entry
rylanmalarchick Jun 18, 2026
ab0dbad
Clarify canonicalize guard name and expand tests
rylanmalarchick Jun 19, 2026
75e5f51
Merge remote-tracking branch 'upstream/main' into extract-canon-disti…
rylanmalarchick Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions doc/releases/changelog-dev.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@

<h3>Improvements 🛠</h3>

* `quantum.extract` canonicalization now looks through a `quantum.insert` at a distinct
static index, rewriting the extract to read from the register feeding the insert. This
removes the false data dependency between wires that act on different qubits of the same
register.
[(#2965)](https://github.com/PennyLaneAI/catalyst/pull/2965)

* The `ResourceAnalysis` pass now reports each loop body and each subroutine as its own entry
instead of folding their gate counts into the caller. Loops with constant bounds appear as `for_loop_<N>`
with their trip count. Loops with dynamic bounds appear as `dyn_for_loop_<N>` with a stable
Expand Down Expand Up @@ -358,6 +364,7 @@ Lillian Frederiksen,
Sengthai Heng,
David Ittah,
Christina Lee,
Rylan Malarchick,
Mehrdad Malekmohammadi,
River McCubbin,
Shuli Shu,
Expand Down
9 changes: 9 additions & 0 deletions mlir/lib/Quantum/IR/QuantumOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,13 +162,22 @@ LogicalResult ExtractOp::canonicalize(ExtractOp extract, mlir::PatternRewriter &
bool staticallyEqual = bothStatic && extract.getIdxAttrAttr() == insert.getIdxAttrAttr();
bool dynamicallyEqual = bothDynamic && extract.getIdx() == insert.getIdx();

bool staticallyDistinct = bothStatic && extract.getIdxAttrAttr() != insert.getIdxAttrAttr();
Comment thread
rylanmalarchick marked this conversation as resolved.
bool inSameBlock = extract->getBlock() == insert->getBlock();

if ((staticallyEqual || dynamicallyEqual) && inSameBlock) {
rewriter.replaceOp(extract, insert.getQubit());
rewriter.replaceOp(insert, insert.getInQreg());
return success();
}

bool insertHasNonExtractUser = llvm::any_of(
insert.getResult().getUsers(), [](Operation *user) { return !isa<ExtractOp>(user); });
if (staticallyDistinct && inSameBlock && insertHasNonExtractUser) {
rewriter.modifyOpInPlace(extract,
[&] { extract.getQregMutable().assign(insert.getInQreg()); });
return success();
}
}
return failure();
}
Expand Down
16 changes: 4 additions & 12 deletions mlir/test/QRef/SemanticConversion/TestFlatCircuits.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,23 +280,15 @@ func.func @test_namedobs_op() -> (!quantum.obs, !quantum.obs) attributes {quantu
// CHECK: [[CNOT:%.+]]:2 = quantum.custom "CNOT"() [[q0]], [[q1]] : !quantum.bit, !quantum.bit
qref.custom "CNOT"() %q0, %q1 : !qref.bit, !qref.bit

// COM: TODO: improve canonicalization patterns to recognize inverse extract-insert pairs where
// COM: inserts are delayed past guaranteed distinct extracts (or vice versa), via statically
// COM: different indices
// CHECK: [[obs_x:%.+]] = quantum.namedobs [[CNOT]]#0[ PauliX] : !quantum.obs
// CHECK: [[insert0:%.+]] = quantum.insert [[qreg]][ 0], [[CNOT]]#0 : !quantum.reg, !quantum.bit
// CHECK: [[insert1:%.+]] = quantum.insert [[insert0]][ 1], [[CNOT]]#1 : !quantum.reg, !quantum.bit

// CHECK: [[extract:%.+]] = quantum.extract [[insert1]][ 0] : !quantum.reg -> !quantum.bit
// CHECK: [[obs_x:%.+]] = quantum.namedobs [[extract]][ PauliX] : !quantum.obs
// CHECK: [[insertX:%.+]] = quantum.insert [[insert1]][ 0], [[extract]] : !quantum.reg, !quantum.bit
%obs_x = qref.namedobs %q0 [ PauliX] : !quantum.obs

// CHECK: [[extract:%.+]] = quantum.extract [[insertX]][ 1] : !quantum.reg -> !quantum.bit
// CHECK: [[obs_z:%.+]] = quantum.namedobs [[extract]][ PauliZ] : !quantum.obs
// CHECK: [[insertZ:%.+]] = quantum.insert [[insertX]][ 1], [[extract]] : !quantum.reg, !quantum.bit
// CHECK: [[obs_z:%.+]] = quantum.namedobs [[CNOT]]#1[ PauliZ] : !quantum.obs
// CHECK: [[insert1:%.+]] = quantum.insert [[insert0]][ 1], [[CNOT]]#1 : !quantum.reg, !quantum.bit
%obs_z = qref.namedobs %q1 [ PauliZ] : !quantum.obs

// CHECK: quantum.dealloc [[insertZ]] : !quantum.reg
// CHECK: quantum.dealloc [[insert1]] : !quantum.reg
Comment thread
mehrdad2m marked this conversation as resolved.
qref.dealloc %a : !qref.reg<2>

// CHECK: return [[obs_x]], [[obs_z]]
Expand Down
6 changes: 2 additions & 4 deletions mlir/test/QRef/SemanticConversion/TestPBC.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,11 +34,9 @@ func.func @test_PPM_op(%angle: f64) -> (i1, i1, i1) attributes {quantum.node} {
// CHECK: [[q1:%.+]] = quantum.extract [[qreg]][ 1] : !quantum.reg -> !quantum.bit
// CHECK: [[m1:%.+]], [[m1_out_qubits:%.+]]:2 = pbc.ppm ["Z", "Y"] [[m0_out_qubit]], [[q1]] : i1, !quantum.bit, !quantum.bit
%m1 = pbc.ref.ppm ["Z", "Y"] %q0, %q1 : i1
// CHECK: [[insert0:%.+]] = quantum.insert [[qreg]][ 0], [[m1_out_qubits]]#0 : !quantum.reg, !quantum.bit
// CHECK: [[insert1:%.+]] = quantum.insert [[insert0]][ 1], [[m1_out_qubits]]#1 : !quantum.reg, !quantum.bit
// CHECK: [[insert1:%.+]] = quantum.insert [[qreg]][ 1], [[m1_out_qubits]]#1 : !quantum.reg, !quantum.bit

// CHECK: [[q0:%.+]] = quantum.extract [[insert1]][ 0] : !quantum.reg -> !quantum.bit
// CHECK: [[m2:%.+]], [[m2_out_qubits:%.+]]:2 = pbc.ppm ["X", "Z"] [[q0]], [[qb]] : i1, !quantum.bit, !quantum.bit
// CHECK: [[m2:%.+]], [[m2_out_qubits:%.+]]:2 = pbc.ppm ["X", "Z"] [[m1_out_qubits]]#0, [[qb]] : i1, !quantum.bit, !quantum.bit
%m2 = pbc.ref.ppm ["X", "Z"] %q0, %qb : i1
// CHECK: [[insert2:%.+]] = quantum.insert [[insert1]][ 0], [[m2_out_qubits]]#0 : !quantum.reg, !quantum.bit

Expand Down
9 changes: 3 additions & 6 deletions mlir/test/QRef/SemanticConversion/TestSubroutines.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,16 +54,13 @@ func.func @main(%arg0: i64, %arg1: f64) -> (!quantum.obs, !quantum.obs) attribut
// CHECK-SAME: (f64, !quantum.bit, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit, !quantum.bit)
func.call @test_extract_before_call(%r2, %r_dyn, %arg1) : (!qref.reg<2>, !qref.reg<?>, f64) -> ()
func.call @test_extract_before_call(%r2, %r_dyn, %arg1) : (!qref.reg<2>, !qref.reg<?>, f64) -> ()
// CHECK: [[insert_20:%.+]] = quantum.insert [[r2]][ 0], [[second_call]]#0 : !quantum.reg, !quantum.bit
// CHECK: [[insert_21:%.+]] = quantum.insert [[insert_20]][ 1], [[second_call]]#1 : !quantum.reg, !quantum.bit
// CHECK: [[insert_21:%.+]] = quantum.insert [[r2]][ 1], [[second_call]]#1 : !quantum.reg, !quantum.bit
// CHECK: [[insert_dyn1:%.+]] = quantum.insert [[r_dyn]][ 1], [[second_call]]#2 : !quantum.reg, !quantum.bit


// CHECK: [[q20:%.+]] = quantum.extract [[insert_21]][ 0] : !quantum.reg -> !quantum.bit
// CHECK: [[obs_q:%.+]] = quantum.compbasis qubits [[q20]] : !quantum.obs
// CHECK: [[obs_q:%.+]] = quantum.compbasis qubits [[second_call]]#0 : !quantum.obs
%q20 = qref.get %r2[0] : !qref.reg<2> -> !qref.bit
%obs_q = qref.compbasis qubits %q20 : !quantum.obs
// CHECK: [[insert_r2:%.+]] = quantum.insert [[insert_21]][ 0], [[q20]] : !quantum.reg, !quantum.bit
// CHECK: [[insert_r2:%.+]] = quantum.insert [[insert_21]][ 0], [[second_call]]#0 : !quantum.reg, !quantum.bit

// CHECK: [[obs_r:%.+]] = quantum.compbasis qreg [[insert_dyn1]] : !quantum.obs
%obs_r = qref.compbasis (qreg %r_dyn : !qref.reg<?>) : !quantum.obs
Expand Down
43 changes: 37 additions & 6 deletions mlir/test/Quantum/CanonicalizationTest.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,15 +71,14 @@ func.func @test_extract_insert_fold(%r1: !quantum.reg, %i: i64) -> !quantum.reg
return %r3 : !quantum.reg
}

// CHECK-LABEL: test_extract_insert_no_fold_static
func.func @test_extract_insert_no_fold_static(%r1: !quantum.reg, %i1: i64, %i2: i64) -> !quantum.reg {
// CHECK: quantum.extract
// CHECK: quantum.insert
// CHECK-LABEL: test_extract_insert_distinct_static_folds
func.func @test_extract_insert_distinct_static_folds(%r1: !quantum.reg, %i1: i64, %i2: i64) -> !quantum.reg {
// CHECK: [[Q:%.+]] = quantum.extract %arg0[ 0]
// CHECK: [[R2:%.+]] = quantum.insert %arg0[ 1], [[Q]]
// CHECK: quantum.insert [[R2]][%arg2], [[Q]]
%q1 = quantum.extract %r1[0] : !quantum.reg -> !quantum.bit
%r2 = quantum.insert %r1[1], %q1 : !quantum.reg, !quantum.bit

// CHECK: quantum.extract
// CHECK: quantum.insert
%q2 = quantum.extract %r2[0] : !quantum.reg -> !quantum.bit
%r3 = quantum.insert %r2[%i1], %q2 : !quantum.reg, !quantum.bit

Expand All @@ -90,6 +89,38 @@ func.func @test_extract_insert_no_fold_static(%r1: !quantum.reg, %i1: i64, %i2:
return %r4 : !quantum.reg
}

// CHECK-LABEL: test_extract_no_redirect_when_insert_is_leaf
func.func @test_extract_no_redirect_when_insert_is_leaf(%r0: !quantum.reg) -> !quantum.bit {
// CHECK: [[Q0:%.+]] = quantum.extract %arg0[ 0]
// CHECK: [[X:%.+]] = quantum.custom "PauliX"() [[Q0]]
// CHECK: [[INS:%.+]] = quantum.insert %arg0[ 0], [[X]]
// CHECK: [[Q1:%.+]] = quantum.extract [[INS]][ 1]
// CHECK: return [[Q1]]
%q0 = quantum.extract %r0[0] : !quantum.reg -> !quantum.bit
%x = quantum.custom "PauliX"() %q0 : !quantum.bit
%r1 = quantum.insert %r0[0], %x : !quantum.reg, !quantum.bit
%q1 = quantum.extract %r1[1] : !quantum.reg -> !quantum.bit
return %q1 : !quantum.bit
}

// CHECK-LABEL: test_extract_through_insert_distinct_index
func.func @test_extract_through_insert_distinct_index(%r0: !quantum.reg) -> !quantum.reg {
// CHECK: [[Q0:%.+]] = quantum.extract %arg0[ 0]
// CHECK: [[X:%.+]] = quantum.custom "PauliX"() [[Q0]]
// CHECK: [[INS0:%.+]] = quantum.insert %arg0[ 0], [[X]]
// CHECK: [[Q1:%.+]] = quantum.extract %arg0[ 1]
// CHECK: [[Z:%.+]] = quantum.custom "PauliZ"() [[Q1]]
// CHECK: [[INS1:%.+]] = quantum.insert [[INS0]][ 1], [[Z]]
// CHECK: return [[INS1]]
%q0 = quantum.extract %r0[0] : !quantum.reg -> !quantum.bit
%x = quantum.custom "PauliX"() %q0 : !quantum.bit
%r1 = quantum.insert %r0[0], %x : !quantum.reg, !quantum.bit
%q1 = quantum.extract %r1[1] : !quantum.reg -> !quantum.bit
%z = quantum.custom "PauliZ"() %q1 : !quantum.bit
%r2 = quantum.insert %r1[1], %z : !quantum.reg, !quantum.bit
return %r2 : !quantum.reg
}
Comment thread
dime10 marked this conversation as resolved.

// CHECK-LABEL: test_extract_insert_constant
func.func @test_extract_insert_constant(%r1: !quantum.reg) -> !quantum.reg {
// CHECK-NOT: arith.constant
Expand Down
4 changes: 3 additions & 1 deletion mlir/test/Quantum/DecomposeLoweringTest.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -749,7 +749,9 @@ module @different_qreg_values{
// CHECK: [[q0:%.+]] = quantum.extract [[reg]][ 0] : !quantum.reg -> !quantum.bit
// CHECK: [[H:%.+]] = quantum.custom "Hadamard"() [[q0]] : !quantum.bit
// CHECK: [[H_insert:%.+]] = quantum.insert [[reg]][ 0], [[H]] : !quantum.reg, !quantum.bit
// CHECK: [[full_insert:%.+]] = quantum.insert [[H_insert]][ 1], [[q1]] : !quantum.reg, !quantum.bit
// CHECK: [[q2_pre:%.+]] = quantum.extract [[reg]][ 2] : !quantum.reg -> !quantum.bit
// CHECK: [[insert_2:%.+]] = quantum.insert [[H_insert]][ 2], [[q2_pre]] : !quantum.reg, !quantum.bit
// CHECK: [[full_insert:%.+]] = quantum.insert [[insert_2]][ 1], [[q1]] : !quantum.reg, !quantum.bit
Comment on lines +752 to +754

@paul0403 paul0403 Jun 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Trying to understand what was adding these new ops. Can you explain a little bit? Because the new pattern should only be changing an operand on an extract op right?

@rylanmalarchick rylanmalarchick Jun 19, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right that the pattern only swaps the extract's qreg operand; it doesn't add ops. The extra extract/insert come from that swap turning off a downstream fold. In the @my_cnot decomposition, qubit 2 gets re-inserted into %3 at index 2. InsertOp::canonicalize only folds an insert of an extracted qubit when the extract and insert share the register (its sameQreg check). Before this change qubit 2 was extract %3[2], so insert %3[2], extract(%3[2]) hit that case and folded away. After the redirect qubit 2 is extract %0[2], so the extract reads %0 while the insert targets %3, sameQreg is false, and the pair stays. It's semantically identical, since the only insert into %3 is at index 0, so %0[2] and %3[2] are the same qubit, and the gate counts match baseline. So it trades the false dependency for two register ops that previously cancelled.

@mehrdad2m mehrdad2m Jun 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that this is semantically equivalent, but it looks slightly awkward. If I understood correctly, this is a case where the pattern ordering is resulting in a different canonicalized form because if say the insert fold was triggered first, it would be folded away right? If that's the case, maybe it is worth having a similar mirrored pattern for the insertOp as well. That could ensure that no matter what is the order of applied canonicalization patterns, we would get the same result.

@rylanmalarchick rylanmalarchick Jun 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah that is the cause of it. The leftover pair only survives because 'InsertOp::canonicalize' folds an inserted-extracted qubit only when the extract and insert name the same register value (its 'sameQreg' check), and the distinct-index rewrite makes the extract read the earlier register while the insert still targets the later one. So the result ends up depending on which pattern the greedy driven happens to apply first.

A mirrored 'InsertOp' pattern is the correct fix in this case. Generalize that 'sameQreg' check from "same register" to "same register up to inserts at distinct static indicies" so 'insert(reg2, extract(reg1, i), i) folds when 'reg1' and 'reg2' agree at index 'i'. That makes the canonical form order-independent, and as a bonus it should (!) fold this awkward extract/insert pair away instead of leaving it, which should also remove the +2-ops drawback noted above.

One coupling I should flag is that this presupposes the rewrite stays in 'canonicalize'. (my open question from my reply last night regardging keeping in 'canonicalize' as the delay-insert form versus moving it to a pass after 'apply-transform-sequence') ; If it is kept in 'canonicalize' I would add the mirrored pattern with its own tests; if it moves to a dedicated pass the confluence concern mostly goes away. Either can be done

@mehrdad2m mehrdad2m Jun 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah that makes sense. I think If making the existing passes' (including the parity-synth pass) assumptions inline with this canonical form then keeping it in the same pass is ideal, but I am not super familiar with the parity-synth implementation to know if it is trivial or not. I'll let @dime10 comment on that.

%0 = quantum.alloc( 3) : !quantum.reg
%1 = quantum.extract %0[ 1] : !quantum.reg -> !quantum.bit
%2 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
Expand Down
Loading