Add hierarchical qubit types

doc/releases/changelog-dev.md

@@ -39,6 +39,11 @@
   into normal PPR and PPMs with SCF dialect to support runtime execution.
   [(#2390)](https://github.com/PennyLaneAI/catalyst/pull/2390)
 
+* New qubit-type specializations have been added Catalyst's MLIR type system. These new qubit types
+  include `!quantum.bit<logical>`, `!quantum.bit<qec>` and `!quantum.bit<physical>`. The original
+  `!quantum.bit` type continues to be supported and used as the default qubit type.
+  [(#2369)](https://github.com/PennyLaneAI/catalyst/pull/2369)
+
 <h3>Documentation 📝</h3>
 
 * Updated the Unified Compiler Cookbook to be compatible with the latest versions of PennyLane and Catalyst.
@@ -47,7 +52,8 @@
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
-Ali Asadi
+Ali Asadi,
+Joey Carter,
 Sengthai Heng,
 Jeffrey Kam,
 Mudit Pandey.

mlir/include/Quantum/IR/QuantumAttrDefs.td

@@ -38,5 +38,28 @@ def NamedObservable : I32EnumAttr<"NamedObservable",
 
 def NamedObservableAttr : EnumAttr<QuantumDialect, NamedObservable, "named_observable">;
 
+def QubitLevel : I32EnumAttr<"QubitLevel",
+    "Qubit levels in the hierarchical qubit representation",
+    [
+        I32EnumAttrCase<"Abstract", 0, "abstract">,
+        I32EnumAttrCase<"Logical",  1, "logical">,
+        I32EnumAttrCase<"QEC",      2, "qec">,
+        I32EnumAttrCase<"Physical", 3, "physical">,
+    ]> {
+    let cppNamespace = "catalyst::quantum";
+    let genSpecializedAttr = 0;
+}
+
+def QubitRole : I32EnumAttr<"QubitRole",
+    "Qubit roles for further specialization in the hierarchical qubit representation",
+    [
+        I32EnumAttrCase<"Null",   0, "null">,
+        I32EnumAttrCase<"Data",   1, "data">,
+        I32EnumAttrCase<"XCheck", 2, "xcheck">,
+        I32EnumAttrCase<"ZCheck", 3, "zcheck">,
+    ]> {
+    let cppNamespace = "catalyst::quantum";
+    let genSpecializedAttr = 0;
+}
 
 #endif // QUANTUM_ATTR_DEFS

mlir/include/Quantum/IR/QuantumInterfaces.td

@@ -79,11 +79,26 @@ def QuantumOperation : OpInterface<"QuantumOperation"> {
 
     let verify = [{
         auto gate = mlir::cast<ConcreteOp>($_op);
+        auto operands = gate.getQubitOperands();
+        auto results = gate.getQubitResults();
 
-        if (gate.getQubitOperands().size() != gate.getQubitResults().size())
+        if (operands.size() != results.size()) {
             return $_op->emitError() <<
-                "number of qubits in input (" << gate.getQubitOperands().size() << ") " <<
-                "and output (" << gate.getQubitResults().size() << ") must be the same";
+                "number of qubits in input (" << operands.size() << ") " <<
+                "and output (" << results.size() << ") must be the same";
+        }
+
+        if (!operands.empty()) {
+            mlir::Type refType = operands[0].getType();
+
+            // Check if all operands and results have the same type as refType
+            auto allMatch = [&](auto value) { return value.getType() == refType; };
+
+            if (!llvm::all_of(operands, allMatch) || !llvm::all_of(results, allMatch)) {
+                return $_op->emitOpError()
+                    << "requires all qubit operands and results to have the same type";
+            }
+        }
 
         return mlir::success();
     }];

mlir/include/Quantum/IR/QuantumOps.td

@@ -346,7 +346,7 @@ class UnitaryGate_Op<string mnemonic, list<Trait> traits = []> :
     let extraClassDeclaration = extraBaseClassDeclaration;
 }
 
-def SetStateOp : Gate_Op<"set_state"> {
+def SetStateOp : Gate_Op<"set_state", [AllTypesMatch<["in_qubits", "out_qubits"]>]> {
     let summary = "Set state to a complex vector.";
     let description = [{
         This operation is useful for simulators implementing state preparation.
@@ -377,7 +377,7 @@ def SetStateOp : Gate_Op<"set_state"> {
 
 }
 
-def SetBasisStateOp : Gate_Op<"set_basis_state"> {
+def SetBasisStateOp : Gate_Op<"set_basis_state", [AllTypesMatch<["in_qubits", "out_qubits"]>]> {
     let summary = "Set basis state.";
     let description = [{
         This operation is useful for simulators implementing set basis state.
@@ -410,7 +410,9 @@ def SetBasisStateOp : Gate_Op<"set_basis_state"> {
 }
 
 def CustomOp : UnitaryGate_Op<"custom", [DifferentiableGate, NoMemoryEffect,
-                                         AttrSizedOperandSegments, AttrSizedResultSegments]> {
+                                         AttrSizedOperandSegments, AttrSizedResultSegments,
+                                         AllTypesMatch<["in_qubits", "out_qubits"]>,
+                                         AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "A generic quantum gate on n qubits with m floating point parameters.";
     let description = [{
     }];
@@ -537,14 +539,16 @@ def CustomOp : UnitaryGate_Op<"custom", [DifferentiableGate, NoMemoryEffect,
 def PauliWord : TypedArrayAttrBase<StrAttr, "A product of Pauli operators, aka a Pauli word.">;
 
 def PauliRotOp : UnitaryGate_Op<"paulirot", [DifferentiableGate, NoMemoryEffect,
-                                         AttrSizedOperandSegments, AttrSizedResultSegments]> {
+                                             AttrSizedOperandSegments, AttrSizedResultSegments,
+                                             AllTypesMatch<["in_qubits", "out_qubits"]>,
+                                             AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "Apply a Pauli Product Rotation";
     let description = [{
         The `quantum.paulirot` operation applies a rotation around a Pauli product
-        operator to the state-vector. 
+        operator to the state-vector.
         The arguments are the rotation angle `angle`, a string representing the
         Pauli product operator, and a set of qubits the operation acts on.
-        Note that this operation is currently not excutable. There isn't a valid 
+        Note that this operation is currently not excutable. There isn't a valid
         lowering path to the LLVM IR.
     }];
 
@@ -565,15 +569,16 @@ def PauliRotOp : UnitaryGate_Op<"paulirot", [DifferentiableGate, NoMemoryEffect,
     let assemblyFormat = [{
         $pauli_product `(` $angle `)` $in_qubits (`adj` $adjoint^)? attr-dict ( `ctrls` `(` $in_ctrl_qubits^ `)` )?  ( `ctrlvals` `(` $in_ctrl_values^ `)` )? `:` type($out_qubits) (`ctrls` type($out_ctrl_qubits)^ )?
     }];
-    
+
     let extraClassDeclaration = extraBaseClassDeclaration # [{
         mlir::ValueRange getAllParams() {
             return getODSOperands(getParamOperandIdx());
         }
     }];
 }
 
-def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOperandSegments]> {
+def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOperandSegments,
+                                              AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "Global Phase.";
 
     let description = [{
@@ -614,7 +619,9 @@ def GlobalPhaseOp : UnitaryGate_Op<"gphase", [DifferentiableGate, AttrSizedOpera
 }
 
 def MultiRZOp : UnitaryGate_Op<"multirz", [DifferentiableGate, NoMemoryEffect,
-                                           AttrSizedOperandSegments, AttrSizedResultSegments]> {
+                                           AttrSizedOperandSegments, AttrSizedResultSegments,
+                                           AllTypesMatch<["in_qubits", "out_qubits"]>,
+                                           AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "Apply an arbitrary multi Z rotation";
     let description = [{
         The `quantum.multirz` operation applies an arbitrary multi Z rotation to the state-vector.
@@ -653,7 +660,9 @@ def MultiRZOp : UnitaryGate_Op<"multirz", [DifferentiableGate, NoMemoryEffect,
 }
 
 def PCPhaseOp : UnitaryGate_Op<"pcphase", [DifferentiableGate, NoMemoryEffect,
-                                           AttrSizedOperandSegments, AttrSizedResultSegments]> {
+                                           AttrSizedOperandSegments, AttrSizedResultSegments,
+                                           AllTypesMatch<["in_qubits", "out_qubits"]>,
+                                           AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "Apply a projector-controlled phase gate";
     let description = [{
         This gate is built from simpler gates like `PhaseShift` and `PauliX` and acts on a group
@@ -703,7 +712,9 @@ def PCPhaseOp : UnitaryGate_Op<"pcphase", [DifferentiableGate, NoMemoryEffect,
 
 
 def QubitUnitaryOp : UnitaryGate_Op<"unitary", [ParametrizedGate, NoMemoryEffect,
-                                                AttrSizedOperandSegments, AttrSizedResultSegments]> {
+                                                AttrSizedOperandSegments, AttrSizedResultSegments,
+                                                AllTypesMatch<["in_qubits", "out_qubits"]>,
+                                                AllTypesMatch<["in_ctrl_qubits", "out_ctrl_qubits"]>]> {
     let summary = "Apply an arbitrary fixed unitary matrix";
     let description = [{
         The `quantum.unitary` operation applies an arbitrary fixed unitary matrix to the
@@ -958,7 +969,7 @@ def HamiltonianOp : Observable_Op<"hamiltonian"> {
 class Measurement_Op<string mnemonic, list<Trait> traits = []> :
         Quantum_Op<mnemonic, traits # [MeasurementProcess]>;
 
-def MeasureOp : Quantum_Op<"measure"> {
+def MeasureOp : Quantum_Op<"measure", [AllTypesMatch<["in_qubit", "out_qubit"]>]> {
     let summary = "A single-qubit projective measurement in the computational basis.";
     let description = [{
     }];

mlir/include/Quantum/IR/QuantumTypes.td

@@ -30,6 +30,29 @@ class Quantum_Type<string name, string typeMnemonic, list<Trait> traits = []>
 
 def QubitType : Quantum_Type<"Qubit", "bit"> {
     let summary = "A value-semantic qubit (state).";
+
+    let parameters = (ins
+        DefaultValuedParameter<"QubitLevel", "QubitLevel::Abstract">:$level,
+        DefaultValuedParameter<"QubitRole", "QubitRole::Null">:$role
+    );
+
+    let assemblyFormat = "(`<` $level^ (`,` $role^ )? `>`)?";
+
+    // This allows the ODS generator to "build" the type automatically using the
+    // 'Abstract' qubit level and 'Null' qubit role and when it needs a default.
+    let builderCall = [{
+        $_builder.getType<quantum::QubitType>(
+            quantum::QubitLevel::Abstract, quantum::QubitRole::Null)
+    }];
+
+    let extraClassDeclaration = [{
+        // Get an instance of the QubitType with the default 'Abstract' level and 'Null' role.
+        static QubitType get(::mlir::MLIRContext *context) {
+            return get(context, quantum::QubitLevel::Abstract, quantum::QubitRole::Null);
+        }
+    }];
+
+    let genVerifyDecl = 1;
 }
 def QuregType : Quantum_Type<"Qureg", "reg"> {
     let summary = "An array of value-semantic qubits (i.e. quantum register).";

mlir/lib/Quantum/IR/CMakeLists.txt

@@ -2,6 +2,7 @@ add_mlir_library(MLIRQuantum
     QuantumDialect.cpp
     QuantumInterfaces.cpp
     QuantumOps.cpp
+    QuantumTypes.cpp
 
     ADDITIONAL_HEADER_DIRS
     ${PROJECT_SOURCE_DIR}/include/Quantum

mlir/lib/Quantum/IR/QuantumTypes.cpp

@@ -0,0 +1,35 @@
+// 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 "mlir/IR/Diagnostics.h"
+#include "mlir/Support/LLVM.h"
+
+#include "Quantum/IR/QuantumAttrDefs.h"
+#include "Quantum/IR/QuantumTypes.h"
+
+using namespace mlir;
+using namespace catalyst::quantum;
+
+LogicalResult QubitType::verify(function_ref<InFlightDiagnostic()> emitError, QubitLevel level,
+                                QubitRole role)
+{
+    // If qubit level is not QEC or Physical, role must be Null
+    // In other words, abstract and logical qubits cannot specify a role
+    if ((level != QubitLevel::QEC && level != QubitLevel::Physical) && role != QubitRole::Null) {
+        return emitError() << "qubit role '" << stringifyQubitRole(role)
+                           << "' is only permitted for qec and physical qubits; "
+                           << "found level '" << stringifyQubitLevel(level) << "'";
+    }
+    return success();
+}