Sara/quantum phase folding

PhaseFoldingTest.mlir

@@ -0,0 +1,116 @@
+func.func @ex_425(%arg0: tensor<2xi64>) -> (!quantum.bit, !quantum.bit) {
+    %reg = quantum.alloc( 2) : !quantum.reg
+    %q0 = quantum.extract %reg[ 0] : !quantum.reg -> !quantum.bit
+    %q1 = quantum.extract %reg[ 1] : !quantum.reg -> !quantum.bit
+
+    // %2 = stablehlo.convert %arg0 : (tensor<2xi64>) -> tensor<2xcomplex<f64>>
+    // %3 = quantum.set_state(%2) %q1 : (tensor<2xcomplex<f64>>, !quantum.bit) -> !quantum.bit
+    // %q3 = quantum.custom "T"() %3 : !quantum.bit
+
+    %tens01 = arith.constant dense<[false]> : tensor<1xi1>
+    %q2 = quantum.set_basis_state(%tens01) %q1 : (tensor<1xi1>, !quantum.bit) -> !quantum.bit
+
+    %q3 = quantum.custom "T"() %q2 : !quantum.bit   // l1
+    %q4:2 = quantum.custom "CNOT"() %q0, %q3 : !quantum.bit, !quantum.bit
+    %q5 = quantum.custom "T"() %q4#0 : !quantum.bit   // l2
+    %65 = quantum.custom "T"() %q4#1 : !quantum.bit   // l3
+
+    func.return %q0, %q3 : !quantum.bit, !quantum.bit
+}
+
+func.func @ex_1() -> (!quantum.bit, !quantum.bit) {
+    %reg = quantum.alloc( 2) : !quantum.reg
+    // %i = arith.constant 1 : index
+    %q0 = quantum.extract %reg[ 0] : !quantum.reg -> !quantum.bit
+    %q1 = quantum.extract %reg[ 1] : !quantum.reg -> !quantum.bit    
+
+    // %c0 = arith.constant 3.141592 : f64
+
+    %tens01 = arith.constant dense<[0]> : tensor<1xi1>
+    %q18 = quantum.set_basis_state(%tens01) %q1 : (tensor<1xi1>, !quantum.bit) -> !quantum.bit
+
+    %q2 = quantum.custom "T"() %q0 : !quantum.bit   // l1
+    // %q2 = quantum.custom "RZ"(%c0) %q0 : !quantum.bit
+    %q3 = quantum.custom "T"() %q18 : !quantum.bit   // l2
+    %q4:2 = quantum.custom "CNOT"() %q2, %q3 : !quantum.bit, !quantum.bit
+    %q5 = quantum.custom "T"() %q4#1 {adjoint}: !quantum.bit    // l3
+    %q6:2 = quantum.custom "CNOT"() %q4#0, %q5 : !quantum.bit, !quantum.bit
+
+    %q7 = quantum.custom "Hadamard"() %q6#1 : !quantum.bit
+
+    %q21 = quantum.custom "PauliX"() %q6#0 : !quantum.bit
+
+    %q8 = quantum.custom "T"() %q21 : !quantum.bit  // l4
+    %q9 = quantum.custom "T"() %q7 : !quantum.bit   // l5
+    %q10:2 = quantum.custom "CNOT"() %q9, %q8 : !quantum.bit, !quantum.bit
+    %q11 = quantum.custom "T"() %q10#1 {adjoint}: !quantum.bit  // l6
+    %q12:2 = quantum.custom "CNOT"() %q10#0, %q11 : !quantum.bit, !quantum.bit
+
+
+    %q13 = quantum.custom "T"() %q12#0 : !quantum.bit   // on q2
+
+    %reg2 = quantum.alloc( 2) : !quantum.reg
+    %p0 = quantum.extract %reg2[ 0] : !quantum.reg -> !quantum.bit
+    %p1 = quantum.extract %reg2[ 1] : !quantum.reg -> !quantum.bit    
+
+    // %c = arith.constant 1.1212 : f64
+    // %q_temp = quantum.custom "RZ"(%c) %q13 : !quantum.bit
+
+    %q14 = quantum.custom "T"() %q13 : !quantum.bit
+    %q15 = quantum.custom "T"() %q14 : !quantum.bit
+    %q16 = quantum.custom "T"() %q15 : !quantum.bit
+    // %q17 = quantum.custom "PauliY"() %q16 : !quantum.bit
+    // %q18 = quantum.custom "T"() %q17 : !quantum.bit
+    // %q19 = quantum.custom "T"() %q18 : !quantum.bit
+    // %q20 = quantum.custom "T"() %q19 : !quantum.bit
+
+    %p2 = quantum.custom "T"() %p0 : !quantum.bit
+    %p3 = quantum.custom "Hadamard"() %p1 : !quantum.bit
+    %p4:2 = quantum.custom "CNOT"() %p3, %p2 : !quantum.bit, !quantum.bit
+    %p5 = quantum.custom "T"() %p4#1 : !quantum.bit
+    %p6:2 = quantum.custom "CNOT"() %p4#0, %p5 : !quantum.bit, !quantum.bit
+    %p7 = quantum.custom "T"() %p6#1 : !quantum.bit
+
+    %q17 = quantum.custom "PauliY"() %q16 : !quantum.bit
+
+    %qb = quantum.alloc_qb : !quantum.bit
+    %h = quantum.custom "PauliX"() %qb : !quantum.bit
+
+    %p8:2 = quantum.custom "SWAP"() %p7, %h : !quantum.bit, !quantum.bit
+    %p9 = quantum.custom "S"() %p8#1 : !quantum.bit
+
+    // %tens01 = arith.constant dense<[false, true]> : tensor<2xi1>
+    // %q18:2 = quantum.set_basis_state(%tens01) %q12#1, %q17 : (tensor<2xi1>, !quantum.bit, !quantum.bit) -> (!quantum.bit, !quantum.bit)
+
+    // func.return %q12#1, %q12#0 : !quantum.bit, !quantum.bit
+    func.return %q12#1, %q12#0 : !quantum.bit, !quantum.bit
+}
+
+// func.func @test() {
+//     %reg = quantum.alloc( 2) : !quantum.reg
+//     %q0 = quantum.extract %reg[ 0] : !quantum.reg -> !quantum.bit
+//     %q1 = quantum.extract %reg[ 1] : !quantum.reg -> !quantum.bit
+
+//     %c0 = arith.constant 0 : index
+//     %c4 = arith.constant 4 : index
+//     %c1 = arith.constant 1 : index
+
+//     %c = arith.constant 1.1212 : f64
+
+//     %q2:2 = quantum.custom "CNOT"() %q0, %q1 : !quantum.bit, !quantum.bit
+//     %q3 = quantum.custom "RZ"(%c) %q2#0 : !quantum.bit
+//     %q4 = quantum.custom "T"() %q2#1 : !quantum.bit
+//     %q5 = quantum.custom "Hadamard"() %q3 : !quantum.bit
+//     %q6 = quantum.custom "PauliX"() %q4 : !quantum.bit
+//     %q7 = quantum.custom "PauliY"() %q5 : !quantum.bit
+//     %q8 = quantum.custom "PauliZ"() %q6 : !quantum.bit
+
+
+//     // %r = scf.for %i = %c0 to %c4 step %c1 iter_args(%x = %q1) -> !quantum.bit {
+//     //     %q1_1 = quantum.custom "PauliX"() %x : !quantum.bit
+//     //     scf.yield %q1_1 : !quantum.bit
+//     // }
+
+
+//     func.return
+// }

main.py

@@ -0,0 +1,61 @@
+import pennylane as qp
+
+
+@qp.qjit(keep_intermediate=2, capture=False)
+@qp.transform(pass_name="phase-folding")
+@qp.qnode(device=qp.device("lightning.qubit", wires=2))
+def circuit_base():
+    qp.CNOT([0, 1])
+    qp.T(1)
+    qp.CNOT([0, 1])
+    qp.CNOT([1, 0])
+    qp.T(0)
+    qp.CNOT([1, 0])
+    return qp.probs()
+
+
+@qp.qjit(keep_intermediate=2, capture=False)
+@qp.transform(pass_name="phase-folding")
+@qp.qnode(device=qp.device("lightning.qubit", wires=2))
+def circuit_ex424():
+    qp.T(0)
+    qp.T(1)
+    qp.CNOT([0, 1])
+    qp.adjoint(qp.T(1))
+    qp.CNOT([0, 1])
+
+    qp.H(1)
+    qp.T(0)
+    qp.T(1)
+    qp.CNOT([1, 0])
+    qp.adjoint(qp.T(0))
+    qp.CNOT([1, 0])
+
+
+# @qp.qjit(keep_intermediate=2, capture=False)
+# # @qp.transform(pass_name="phase-folding")
+# @qp.qnode(device=qp.device("lightning.qubit", wires=2))
+# def circuit_ex425():
+#     # qp.StatePrep([0, 1], wires=1)
+#     qp.BasisState(0, wires=1)   # becomes dynamic, so not yet
+#     qp.T(1)
+#     qp.CNOT([0, 1])
+#     qp.T(0)
+#     qp.T(1)
+
+
+@qp.qjit(keep_intermediate=2, capture=False)
+@qp.transform(pass_name="phase-folding")
+@qp.qnode(device=qp.device("lightning.qubit", wires=1))
+def circuit_ex426():
+    qp.T(0)
+    qp.PauliX(0)
+    qp.adjoint(qp.T(0))
+    qp.PauliX(0)
+    return qp.probs()
+
+
+circuit_base()
+circuit_ex424()
+# circuit_ex425()
+circuit_ex426()

mlir/include/Quantum/Transforms/Passes.td

@@ -218,6 +218,12 @@ def CombineGlobalPhasesPass : Pass<"combine-global-phases", "mlir::ModuleOp"> {
     let summary = "Scope-wise merging of global phase operators (safe).";
 }
 
+def PhaseFoldingPass : Pass<"phase-folding"> {
+    let summary = "Perform phase-folding optimization on pure quantum blocks to reduce T-count.";
+
+    // let constructor = "catalyst::createMyHelloWorldPass()";
+}
+
 // ----- Quantum circuit transformation passes end ----- //
 
 #endif // QUANTUM_PASSES

mlir/lib/Quantum/Transforms/CMakeLists.txt

@@ -31,6 +31,18 @@ file(GLOB SRC
     loop_boundary_optimization.cpp
     LoopBoundaryOptimizationPatterns.cpp
     reference_semantics_conversion.cpp
+    phase_folding.cpp
+    SymbolicAnalysis/PhaseBucket.h
+    SymbolicAnalysis/PhaseBucket.cpp
+    SymbolicAnalysis/Parity.h
+    SymbolicAnalysis/Parity.cpp
+    SymbolicAnalysis/PhasePolynomial.h
+    SymbolicAnalysis/PhasePolynomial.cpp
+    SymbolicAnalysis/AffineTransform.h
+    SymbolicAnalysis/AffineTransform.cpp
+    SymbolicAnalysis/SymbolicCircuit.h
+    SymbolicAnalysis/SymbolicCircuit.cpp
+    SymbolicAnalysis/Gate.h
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)

mlir/lib/Quantum/Transforms/SymbolicAnalysis/AffineTransform.cpp

@@ -0,0 +1,87 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "AffineTransform.h"
+
+/*
+    Constructors:
+*/
+AffineTransform::AffineTransform(size_t n)
+{
+    exprMatrix.reserve(n);
+    for (size_t i = 0; i < n; i++) {
+        exprMatrix.push_back(Parity::eVec(n, i + 1));
+    }
+}
+
+AffineTransform AffineTransform::identity(size_t n) { return AffineTransform(n); }
+
+/*
+    Operators:
+*/
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const AffineTransform &trans)
+{
+    for (auto it = trans.exprMatrix.begin(); it != trans.exprMatrix.end(); it++) {
+        os << *it << '\n';
+    }
+    return os;
+}
+
+std::string AffineTransform::algebraicView(size_t qubitNum) const
+{
+    std::string res = "";
+    for (size_t i = 0; i < exprMatrix.size(); i++) {
+        res +=
+            ("x'" + std::to_string(i + 1) + " = " + exprMatrix[i].algebraicView(qubitNum) + '\n');
+    }
+    return res;
+}
+
+/*
+    Getters and Setters:
+*/
+void AffineTransform::setRow(size_t row, const Parity &parity) { getRowMutable(row) = parity; }
+
+void AffineTransform::resetRow(size_t row) { getRowMutable(row).reset(); }
+
+void AffineTransform::flipAffineValueAtRow(size_t row) { getRowMutable(row).flipAffineValue(); }
+
+/*
+    Methods:
+*/
+void AffineTransform::addRowWithParity(size_t row, const Parity &parity)
+{
+    getRowMutable(row) += parity;
+}
+
+void AffineTransform::addRows(size_t sourceRow, size_t targetRow)
+{ // E_i,j
+    getRowMutable(targetRow) += getRow(sourceRow);
+}
+
+void AffineTransform::swapRows(size_t row1, size_t row2)
+{
+    std::swap(getRowMutable(row1), getRowMutable(row2));
+}
+
+void AffineTransform::extendTo(size_t newRowNum, size_t auxVarNum)
+{
+    if (newRowNum > getRowNum()) {
+        exprMatrix.reserve(newRowNum);
+        for (size_t i = getRowNum() + 1; i <= newRowNum; i++) {
+            exprMatrix.push_back(Parity::eVec(newRowNum + auxVarNum, i + auxVarNum));
+        }
+    } // it might be more efficient to resize to newRowNum, and then just turn the eVec bits on for
+      // the new rows.
+}

mlir/lib/Quantum/Transforms/SymbolicAnalysis/AffineTransform.h

@@ -0,0 +1,76 @@
+// Copyright 2026 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cassert>
+#include <string_view>
+#include <utility>
+#include <vector>
+
+#include "llvm/Support/raw_ostream.h"
+
+#include "Parity.h"
+
+class AffineTransform {
+  public:
+    // Constructors
+    AffineTransform() = default;
+    AffineTransform(const Parity *rows, size_t n) : exprMatrix(rows, rows + n) {}
+
+    // Static Factories
+    static AffineTransform identity(size_t n);
+
+    // Operators
+    friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const AffineTransform &trans);
+    std::string algebraicView(size_t qubitNum) const;
+
+    // Getters
+    [[nodiscard]] size_t getRowNum() const;
+    [[nodiscard]] size_t getColNum(size_t row) const;
+    [[nodiscard]] const Parity &getRow(size_t row) const;
+    [[nodiscard]] Parity &getRowMutable(size_t row) const;
+
+    // Setters
+    void setRow(size_t row, const Parity &parity);
+    void resetRow(size_t row);
+    void flipAffineValueAtRow(size_t row);
+
+    // Methods
+    void extendTo(size_t newRowNum, size_t auxVarNum);
+    void swapRows(size_t row1, size_t row2);
+    void addRows(size_t sourceRow, size_t targetRow);
+    void addRowWithParity(size_t row, const Parity &parity);
+
+  private:
+    std::vector<Parity> exprMatrix; // n x (n + m + 1) binary matrix
+
+    // Helper Methods
+    explicit AffineTransform(size_t n); // Identity matrix by default
+};
+
+inline size_t AffineTransform::getRowNum() const { return exprMatrix.size(); }
+
+inline size_t AffineTransform::getColNum(size_t row) const { return getRow(row).getVarNum(); }
+
+inline const Parity &AffineTransform::getRow(size_t row) const
+{
+    assert(row >= 0 && row < getRowNum());
+    return exprMatrix[row];
+}
+
+inline Parity &AffineTransform::getRowMutable(size_t row) const
+{
+    return const_cast<Parity &>(static_cast<const AffineTransform &>(*this).getRow(row));
+}