Add canonicalize to extract (#2162)

**Context:**

**Description of the Change:**

Fix canonicalization of eliminating redundant `quantum.insert` and
`quantum.extract` pairs. When extracting a qubit immediately after
inserting it at the same index, the operations can be cancelled out
while properly updating remaining uses of the register. For an example:
```mlir
// Before canonicalization
%1 = quantum.insert %0[%idx], %qubit1 : !quantum.reg, !quantum.bit
%2 = quantum.extract %1[%idx] : !quantum.reg -> !quantum.bit
...
%3 = quantum.insert %1[%i0], %qubit2 : !quantum.reg, !quantum.bit
%4 = quantum.extract %1[%i1] : !quantum.reg -> !quantum.bit
// ... use %1
// ... use %4

// After canonicalization
// %2 directly uses %qubit1
// %3 and %4 updated to use %0 instead of %1
%3 = quantum.insert %0[%i0], %qubit2 : !quantum.reg, !quantum.bit
%4 = quantum.extract %0[%i1] : !quantum.reg -> !quantum.bit
// ... use %qubit1
// ... use %4
```

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

Author: rniczh

doc/releases/changelog-dev.md

@@ -44,6 +44,31 @@
 
 * Fixes the translation of a workflow with different transforms applied to different qnodes.
   [(#2167)](https://github.com/PennyLaneAI/catalyst/pull/2167)
+  
+* Fix canonicalization of eliminating redundant `quantum.insert` and `quantum.extract` pairs.
+  When extracting a qubit immediately after inserting it at the same index, the operations can
+  be cancelled out while properly updating remaining uses of the register.
+  [(#2162)](https://github.com/PennyLaneAI/catalyst/pull/2162)
+  For an example:
+```mlir
+// Before canonicalization
+%1 = quantum.insert %0[%idx], %qubit1 : !quantum.reg, !quantum.bit
+%2 = quantum.extract %1[%idx] : !quantum.reg -> !quantum.bit
+...
+%3 = quantum.insert %1[%i0], %qubit2 : !quantum.reg, !quantum.bit
+%4 = quantum.extract %1[%i1] : !quantum.reg -> !quantum.bit
+// ... use %1
+// ... use %4
+
+// After canonicalization
+// %2 directly uses %qubit1
+// %3 and %4 updated to use %0 instead of %1
+%3 = quantum.insert %0[%i0], %qubit2 : !quantum.reg, !quantum.bit
+%4 = quantum.extract %0[%i1] : !quantum.reg -> !quantum.bit
+// ... use %qubit1
+// ... use %4
+```
+
 
 <h3>Internal changes ⚙️</h3>
 
@@ -75,4 +100,5 @@ Ali Asadi,
 Christina Lee,
 River McCubbin,
 Roberto Turrado,
-Paul Haochen Wang.
+Paul Haochen Wang,
+Hongsheng Zheng.

mlir/lib/Quantum/IR/QuantumOps.cpp

@@ -155,11 +155,10 @@ LogicalResult ExtractOp::canonicalize(ExtractOp extract, mlir::PatternRewriter &
         bool bothDynamic = !extract.getIdxAttr().has_value() && !insert.getIdxAttr().has_value();
         bool staticallyEqual = bothStatic && extract.getIdxAttrAttr() == insert.getIdxAttrAttr();
         bool dynamicallyEqual = bothDynamic && extract.getIdx() == insert.getIdx();
-        // if other users of insert are also `insert`, we are good to go
-        bool valid = llvm::all_of(insert.getResult().getUsers(), [&](Operation *op) {
-            return isa<InsertOp>(op) || op == extract.getOperation();
-        });
-        if ((staticallyEqual || dynamicallyEqual) && valid) {
+
+        bool inSameBlock = extract->getBlock() == insert->getBlock();
+
+        if ((staticallyEqual || dynamicallyEqual) && inSameBlock) {
             rewriter.replaceOp(extract, insert.getQubit());
             rewriter.replaceOp(insert, insert.getInQreg());
             return success();

mlir/test/Quantum/DecomposeLoweringTest.mlir

@@ -388,12 +388,10 @@ module @multi_wire_cnot_decomposition {
     // CHECK: [[QUBIT1:%.+]] = quantum.extract [[REG]][[[EXTRACTED]]] : !quantum.reg -> !quantum.bit
     // CHECK: [[RZ1:%.+]] = quantum.custom "RZ"([[CST_PI]]) [[QUBIT1]] : !quantum.bit
     // CHECK: [[RY1:%.+]] = quantum.custom "RY"([[CST_PI2]]) [[RZ1]] : !quantum.bit
-    // CHECK: [[INSERT1:%.+]] = quantum.insert [[REG]][[[EXTRACTED]]], [[RY1]] : !quantum.reg, !quantum.bit
     // CHECK: [[EXTRACTED2:%.+]] = tensor.extract [[RESHAPE1]][] : tensor<i64>
-    // CHECK: [[QUBIT0:%.+]] = quantum.extract [[INSERT1]][[[EXTRACTED2]]] : !quantum.reg -> !quantum.bit
-    // CHECK: [[QUBIT1_2:%.+]] = quantum.extract [[INSERT1]][[[EXTRACTED]]] : !quantum.reg -> !quantum.bit
-    // CHECK: [[CZ_RESULT:%.+]]:2 = quantum.custom "CZ"() [[QUBIT0]], [[QUBIT1_2]] : !quantum.bit, !quantum.bit
-    // CHECK: [[INSERT2:%.+]] = quantum.insert [[INSERT1]][[[EXTRACTED2]]], [[CZ_RESULT]]#0 : !quantum.reg, !quantum.bit
+    // CHECK: [[QUBIT0:%.+]] = quantum.extract [[REG]][[[EXTRACTED2]]] : !quantum.reg -> !quantum.bit
+    // CHECK: [[CZ_RESULT:%.+]]:2 = quantum.custom "CZ"() [[QUBIT0]], [[RY1]] : !quantum.bit, !quantum.bit
+    // CHECK: [[INSERT2:%.+]] = quantum.insert [[REG]][[[EXTRACTED2]]], [[CZ_RESULT]]#0 : !quantum.reg, !quantum.bit
     // CHECK: [[RZ2:%.+]] = quantum.custom "RZ"([[CST_PI]]) [[CZ_RESULT]]#1 : !quantum.bit
     // CHECK: [[RY2:%.+]] = quantum.custom "RY"([[CST_PI2]]) [[RZ2]] : !quantum.bit
     // CHECK: [[INSERT3:%.+]] = quantum.insert [[INSERT2]][[[EXTRACTED]]], [[RY2]] : !quantum.reg, !quantum.bit