bloqade-circuit 0.7.13__py3-none-any.whl → 0.8.0__py3-none-any.whl

bloqade/cirq_utils/lowering.py

@@ -0,0 +1,664 @@
+from typing import Any
+from dataclasses import field, dataclass
+
+import cirq
+from kirin import ir, types, lowering
+from kirin.rewrite import Walk, CFGCompactify
+from kirin.dialects import py, scf, func, ilist
+
+from bloqade import qubit
+from bloqade.squin import gate, noise, kernel, qalloc
+
+
+def load_circuit(
+    circuit: cirq.Circuit,
+    kernel_name: str = "main",
+    dialects: ir.DialectGroup = kernel,
+    register_as_argument: bool = False,
+    return_register: bool = False,
+    register_argument_name: str = "q",
+    globals: dict[str, Any] | None = None,
+    file: str | None = None,
+    lineno_offset: int = 0,
+    col_offset: int = 0,
+    compactify: bool = True,
+):
+    """Converts a cirq.Circuit object into a squin kernel.
+
+    Args:
+        circuit (cirq.Circuit): The circuit to load.
+
+    Keyword Args:
+        kernel_name (str): The name of the kernel to load. Defaults to "main".
+        dialects (ir.DialectGroup | None): The dialects to use. Defaults to `squin.kernel`.
+        register_as_argument (bool): Determine whether the resulting kernel function should accept
+            a single `ilist.IList[Qubit, Any]` argument that is a list of qubits used within the
+            function. This allows you to compose kernel functions generated from circuits.
+            Defaults to `False`.
+        return_register (bool): Determine whether the resulting kernel functionr returns a
+            single value of type `ilist.IList[Qubit, Any]` that is the list of qubits used
+            in the kernel function. Useful when you want to compose multiple kernel functions
+            generated from circuits. Defaults to `False`.
+        register_argument_name (str): The name of the argument that represents the qubit register.
+            Only used when `register_as_argument=True`. Defaults to "q".
+        globals (dict[str, Any] | None): The global variables to use. Defaults to None.
+        file (str | None): The file name for error reporting. Defaults to None.
+        lineno_offset (int): The line number offset for error reporting. Defaults to 0.
+        col_offset (int): The column number offset for error reporting. Defaults to 0.
+        compactify (bool): Whether to compactify the output. Defaults to True.
+
+    ## Usage Examples:
+
+    ```python
+    # from cirq's "hello qubit" example
+    import cirq
+    from bloqade.cirq_utils import load_circuit
+
+    # Pick a qubit.
+    qubit = cirq.GridQubit(0, 0)
+
+    # Create a circuit.
+    circuit = cirq.Circuit(
+        cirq.X(qubit)**0.5,  # Square root of NOT.
+        cirq.measure(qubit, key='m')  # Measurement.
+    )
+
+    # load the circuit as squin
+    main = load_circuit(circuit)
+
+    # print the resulting IR
+    main.print()
+    ```
+
+    You can also compose kernel functions generated from circuits by passing in
+    and / or returning the respective quantum registers:
+
+    ```python
+    import cirq
+    from bloqade.cirq_utils import load_circuit
+    from bloqade import squin
+
+    q = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(cirq.H(q[0]), cirq.CX(*q))
+
+    get_entangled_qubits = load_circuit(
+        circuit, return_register=True, kernel_name="get_entangled_qubits"
+    )
+    get_entangled_qubits.print()
+
+    entangle_qubits = load_circuit(
+        circuit, register_as_argument=True, kernel_name="entangle_qubits"
+    )
+
+    @squin.kernel
+    def main():
+        qreg = get_entangled_qubits()
+        qreg2 = squin.qalloc(1)
+        entangle_qubits([qreg[1], qreg2[0]])
+        return squin.qubit.measure(qreg2)
+    ```
+    """
+
+    target = Squin(dialects=dialects, circuit=circuit)
+    body = target.run(
+        circuit,
+        source=str(circuit),  # TODO: proper source string
+        file=file,
+        globals=globals,
+        lineno_offset=lineno_offset,
+        col_offset=col_offset,
+        compactify=compactify,
+        register_as_argument=register_as_argument,
+        register_argument_name=register_argument_name,
+    )
+
+    if return_register:
+        return_value = target.qreg
+    else:
+        return_value = func.ConstantNone()
+        body.blocks[0].stmts.append(return_value)
+
+    return_node = func.Return(value_or_stmt=return_value)
+    body.blocks[0].stmts.append(return_node)
+
+    self_arg_name = kernel_name + "_self"
+    arg_names = [self_arg_name]
+    if register_as_argument:
+        args = (target.qreg.type,)
+        arg_names.append(register_argument_name)
+    else:
+        args = ()
+
+    # NOTE: add _self as argument; need to know signature before so do it after lowering
+    signature = func.Signature(args, return_node.value.type)
+    body.blocks[0].args.insert_from(
+        0,
+        types.Generic(ir.Method, types.Tuple.where(signature.inputs), signature.output),
+        self_arg_name,
+    )
+
+    code = func.Function(
+        sym_name=kernel_name,
+        signature=signature,
+        body=body,
+    )
+
+    mt = ir.Method(
+        mod=None,
+        py_func=None,
+        sym_name=kernel_name,
+        arg_names=arg_names,
+        dialects=dialects,
+        code=code,
+    )
+
+    assert (run_pass := kernel.run_pass) is not None
+    run_pass(mt, typeinfer=True)
+
+    return mt
+
+
+CirqNode = (
+    cirq.Circuit
+    | cirq.FrozenCircuit
+    | cirq.Moment
+    | cirq.Gate
+    | cirq.Qid
+    | cirq.Operation
+)
+
+DecomposeNode = (
+    cirq.SwapPowGate
+    | cirq.ISwapPowGate
+    | cirq.PhasedXPowGate
+    | cirq.PhasedXZGate
+    | cirq.CSwapGate
+    | cirq.XXPowGate
+    | cirq.YYPowGate
+    | cirq.CCXPowGate
+    | cirq.CCZPowGate
+)
+
+
+@dataclass
+class Squin(lowering.LoweringABC[cirq.Circuit]):
+    """Lower a cirq.Circuit object to a squin kernel"""
+
+    circuit: cirq.Circuit
+    qreg: ir.SSAValue = field(init=False)
+    qreg_index: dict[cirq.Qid, int] = field(init=False, default_factory=dict)
+    next_qreg_index: int = field(init=False, default=0)
+
+    two_qubit_paulis = (
+        "IX",
+        "IY",
+        "IZ",
+        "XI",
+        "XX",
+        "XY",
+        "XZ",
+        "YI",
+        "YX",
+        "YY",
+        "YZ",
+        "ZI",
+        "ZX",
+        "ZY",
+        "ZZ",
+    )
+
+    def __post_init__(self):
+        # TODO: sort by cirq ordering
+        qbits = sorted(self.circuit.all_qubits())
+        self.qreg_index = {qid: idx for (idx, qid) in enumerate(qbits)}
+
+    def lower_qubit_getindex(self, state: lowering.State[cirq.Circuit], qid: cirq.Qid):
+        index = self.qreg_index[qid]
+        index_ssa = state.current_frame.push(py.Constant(index)).result
+        qbit_getitem = state.current_frame.push(py.GetItem(self.qreg, index_ssa))
+        return qbit_getitem.result
+
+    def lower_qubit_getindices(
+        self, state: lowering.State[cirq.Circuit], qids: tuple[cirq.Qid, ...]
+    ):
+        qbits_getitem = [self.lower_qubit_getindex(state, qid) for qid in qids]
+        qbits = state.current_frame.push(ilist.New(values=qbits_getitem))
+        return qbits.result
+
+    def run(
+        self,
+        stmt: cirq.Circuit,
+        *,
+        source: str | None = None,
+        globals: dict[str, Any] | None = None,
+        file: str | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+        register_as_argument: bool = False,
+        register_argument_name: str = "q",
+    ) -> ir.Region:
+
+        state = lowering.State(
+            self,
+            file=file,
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+        )
+
+        with state.frame([stmt], globals=globals, finalize_next=False) as frame:
+
+            # NOTE: need a register of qubits before lowering statements
+            if register_as_argument:
+                # NOTE: register as argument to the kernel; we have freedom of choice for the name here
+                frame.curr_block.args.append_from(
+                    ilist.IListType[qubit.QubitType, types.Any],
+                    name=register_argument_name,
+                )
+                self.qreg = frame.curr_block.args[0]
+            else:
+                # NOTE: create a new register of appropriate size
+                n_qubits = len(self.qreg_index)
+                n = frame.push(py.Constant(n_qubits))
+                self.qreg = frame.push(
+                    func.Invoke((n.result,), callee=qalloc, kwargs=())
+                ).result
+
+            self.visit(state, stmt)
+
+            if compactify:
+                Walk(CFGCompactify()).rewrite(frame.curr_region)
+
+            region = frame.curr_region
+
+        return region
+
+    def visit(
+        self, state: lowering.State[cirq.Circuit], node: CirqNode
+    ) -> lowering.Result:
+        name = node.__class__.__name__
+        return getattr(self, f"visit_{name}", self.generic_visit)(state, node)
+
+    def generic_visit(self, state: lowering.State[cirq.Circuit], node: CirqNode):
+        if isinstance(node, CirqNode):
+            raise lowering.BuildError(
+                f"Cannot lower {node.__class__.__name__} node: {node}"
+            )
+        raise lowering.BuildError(f"Cannot lower {node}")
+
+        # return self.visit_Operation(state, node)
+
+    def lower_literal(self, state: lowering.State[cirq.Circuit], value) -> ir.SSAValue:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def lower_global(
+        self, state: lowering.State[cirq.Circuit], node: CirqNode
+    ) -> lowering.LoweringABC.Result:
+        raise lowering.BuildError("Literals not supported in cirq circuit")
+
+    def visit_Circuit(
+        self,
+        state: lowering.State[cirq.Circuit],
+        node: cirq.Circuit | cirq.FrozenCircuit,
+    ) -> lowering.Result:
+        for moment in node:
+            self.visit_Moment(state, moment)
+
+    def visit_Moment(
+        self, state: lowering.State[cirq.Circuit], node: cirq.Moment
+    ) -> lowering.Result:
+        for op_ in node.operations:
+            self.visit(state, op_)
+
+    def visit_GateOperation(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        if isinstance(node.gate, DecomposeNode):
+            # NOTE: easier to decompose these, but for that we need the qubits too,
+            # so we need to do this within this method
+            for subnode in cirq.decompose_once(node):
+                self.visit(state, subnode)
+            return
+
+        # NOTE: just forward to the appropriate method by getting the name
+        name = node.gate.__class__.__name__
+        return getattr(self, f"visit_{name}", self.generic_visit)(state, node)
+
+    def visit_TaggedOperation(
+        self, state: lowering.State[cirq.Circuit], node: cirq.TaggedOperation
+    ):
+        return self.visit(state, node.untagged)
+
+    def visit_ClassicallyControlledOperation(
+        self,
+        state: lowering.State[cirq.Circuit],
+        node: cirq.ClassicallyControlledOperation,
+    ):
+        conditions: list[ir.SSAValue] = []
+        for outcome in node.classical_controls:
+            key = outcome.key
+            if isinstance(key, cirq.MeasurementKey):
+                key = key.name
+            measurement_outcome = state.current_frame.defs[key]
+
+            if measurement_outcome.type.is_subseteq(ilist.IListType):
+                # NOTE: there is currently no convenient ilist.any method, so we need to use foldl
+                # with a simple function that just does an or
+
+                def bool_op_or(x: bool, y: bool) -> bool:
+                    return x or y
+
+                f_code = state.current_frame.push(
+                    lowering.Python(self.dialects).python_function(bool_op_or)
+                )
+                fn = ir.Method(
+                    mod=None,
+                    py_func=bool_op_or,
+                    sym_name="bool_op_or",
+                    arg_names=[],
+                    dialects=self.dialects,
+                    code=f_code,
+                )
+                f_const = state.current_frame.push(py.constant.Constant(fn))
+                init_val = state.current_frame.push(py.Constant(False)).result
+                condition = state.current_frame.push(
+                    ilist.Foldl(f_const.result, measurement_outcome, init=init_val)
+                ).result
+            else:
+                condition = measurement_outcome
+
+            conditions.append(condition)
+
+        if len(conditions) == 1:
+            condition = conditions[0]
+        else:
+            condition = state.current_frame.push(
+                py.boolop.And(conditions[0], conditions[1])
+            ).result
+            for next_cond in conditions[2:]:
+                condition = state.current_frame.push(
+                    py.boolop.And(condition, next_cond)
+                ).result
+
+        then_stmt = self.visit(state, node.without_classical_controls())
+
+        assert isinstance(
+            then_stmt, ir.Statement
+        ), f"Expected operation of classically controlled node {node} to be lowered to a statement, got type {type(then_stmt)}. \
+        Please report this issue!"
+
+        # NOTE: remove stmt from parent block
+        then_stmt.detach()
+        then_body = ir.Block((then_stmt,))
+        then_body.args.append_from(types.Bool, name="cond")
+        then_body.stmts.append(scf.Yield())
+
+        else_body = ir.Block(())
+        else_body.args.append_from(types.Bool, name="cond")
+        else_body.stmts.append(scf.Yield())
+
+        return state.current_frame.push(
+            scf.IfElse(condition, then_body=then_body, else_body=else_body)
+        )
+
+    def visit_MeasurementGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        qubits = self.lower_qubit_getindices(state, node.qubits)
+        stmt = state.current_frame.push(qubit.stmts.Measure(qubits))
+
+        # NOTE: add for classically controlled lowering
+        key = node.gate.key
+        if isinstance(key, cirq.MeasurementKey):
+            key = key.name
+        state.current_frame.defs[key] = stmt.result
+
+        return stmt
+
+    def visit_SingleQubitPauliStringGateOperation(
+        self,
+        state: lowering.State[cirq.Circuit],
+        node: cirq.SingleQubitPauliStringGateOperation,
+    ):
+        if isinstance(node.pauli, cirq.IdentityGate):
+            # TODO: do we need an identity gate in gate?
+            return
+
+        qargs = self.lower_qubit_getindices(state, (node.qubit,))
+        match node.pauli:
+            case cirq.X:
+                gate_stmt = gate.stmts.X
+            case cirq.Y:
+                gate_stmt = gate.stmts.Y
+            case cirq.Z:
+                gate_stmt = gate.stmts.Z
+            case _:
+                raise lowering.BuildError(f"Unexpected Pauli operation {node.pauli}")
+
+        return state.current_frame.push(gate_stmt(qargs))
+
+    def visit_HPowGate(self, state: lowering.State[cirq.Circuit], node: cirq.HPowGate):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+
+        if node.gate.exponent % 2 == 1:
+            return state.current_frame.push(gate.stmts.H(qargs))
+
+        # NOTE: decompose into products of paulis for arbitrary exponents according to _decompose_ method
+        for subnode in cirq.decompose_once(node):
+            self.visit(state, subnode)
+
+    def visit_XPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+        if node.gate.exponent % 2 == 1:
+            return state.current_frame.push(gate.stmts.X(qargs))
+
+        angle = state.current_frame.push(py.Constant(0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Rx(angle.result, qargs))
+
+    def visit_YPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+        if node.gate.exponent % 2 == 1:
+            return state.current_frame.push(gate.stmts.Y(qargs))
+
+        angle = state.current_frame.push(py.Constant(0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Ry(angle.result, qargs))
+
+    def visit_ZPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+
+        if abs(node.gate.exponent) == 0.5:
+            adjoint = node.gate.exponent < 0
+            return state.current_frame.push(gate.stmts.S(adjoint=adjoint, qubits=qargs))
+
+        if abs(node.gate.exponent) == 0.25:
+            adjoint = node.gate.exponent < 0
+            return state.current_frame.push(gate.stmts.T(adjoint=adjoint, qubits=qargs))
+
+        if node.gate.exponent % 2 == 1:
+            return state.current_frame.push(gate.stmts.Z(qubits=qargs))
+
+        angle = state.current_frame.push(py.Constant(0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Rz(angle.result, qargs))
+
+    def visit_Rx(self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+        angle = state.current_frame.push(py.Constant(value=0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Rx(angle.result, qargs))
+
+    def visit_Ry(self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+        angle = state.current_frame.push(py.Constant(value=0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Ry(angle.result, qargs))
+
+    def visit_Rz(self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation):
+        qargs = self.lower_qubit_getindices(state, node.qubits)
+        angle = state.current_frame.push(py.Constant(value=0.5 * node.gate.exponent))
+        return state.current_frame.push(gate.stmts.Rz(angle.result, qargs))
+
+    def visit_CXPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        if node.gate.exponent % 2 == 0:
+            return
+
+        if node.gate.exponent % 2 != 1:
+            raise lowering.BuildError("Exponents of CX gate are not supported!")
+
+        control, target = node.qubits
+        control_qarg = self.lower_qubit_getindices(state, (control,))
+        target_qarg = self.lower_qubit_getindices(state, (target,))
+        return state.current_frame.push(
+            gate.stmts.CX(controls=control_qarg, targets=target_qarg)
+        )
+
+    def visit_CZPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        if node.gate.exponent % 2 == 0:
+            return
+
+        if node.gate.exponent % 2 != 1:
+            raise lowering.BuildError("Exponents of CZ gate are not supported!")
+
+        control, target = node.qubits
+        control_qarg = self.lower_qubit_getindices(state, (control,))
+        target_qarg = self.lower_qubit_getindices(state, (target,))
+        return state.current_frame.push(
+            gate.stmts.CZ(controls=control_qarg, targets=target_qarg)
+        )
+
+    def visit_ZZPowGate(
+        self, state: lowering.State[cirq.Circuit], node: cirq.GateOperation
+    ):
+        if node.gate.exponent % 2 == 0:
+            return
+
+        qubit1, qubit2 = node.qubits
+        qarg1 = self.lower_qubit_getindices(state, (qubit1,))
+        qarg2 = self.lower_qubit_getindices(state, (qubit2,))
+
+        if node.gate.exponent % 2 == 1:
+            state.current_frame.push(gate.stmts.X(qarg1))
+            state.current_frame.push(gate.stmts.X(qarg2))
+            return
+
+        # NOTE: arbitrary exponent, write as CX * Rz * CX (up to global phase)
+        state.current_frame.push(gate.stmts.CX(qarg1, qarg2))
+        angle = state.current_frame.push(py.Constant(0.5 * node.gate.exponent))
+        state.current_frame.push(gate.stmts.Rz(angle.result, qarg2))
+        state.current_frame.push(gate.stmts.CX(qarg1, qarg2))
+
+    def visit_ControlledOperation(
+        self, state: lowering.State[cirq.Circuit], node: cirq.ControlledOperation
+    ):
+        match node.gate.sub_gate:
+            case cirq.X:
+                stmt = gate.stmts.CX
+            case cirq.Y:
+                stmt = gate.stmts.CY
+            case cirq.Z:
+                stmt = gate.stmts.CZ
+            case _:
+                raise lowering.BuildError(
+                    f"Cannot lowering controlled operation: {node}"
+                )
+
+        control, target = node.qubits
+        control_qarg = self.lower_qubit_getindices(state, (control,))
+        target_qarg = self.lower_qubit_getindices(state, (target,))
+        return state.current_frame.push(stmt(control_qarg, target_qarg))
+
+    def visit_FrozenCircuit(
+        self, state: lowering.State[cirq.Circuit], node: cirq.FrozenCircuit
+    ):
+        return self.visit_Circuit(state, node)
+
+    def visit_CircuitOperation(
+        self, state: lowering.State[cirq.Circuit], node: cirq.CircuitOperation
+    ):
+        reps = node.repetitions
+
+        if not isinstance(reps, int):
+            raise lowering.BuildError(
+                f"Cannot lower CircuitOperation with non-integer repetitions: {node}"
+            )
+
+        if reps > 1:
+            raise lowering.BuildError(
+                "Repetitions of circuit operatiosn not yet supported"
+            )
+
+        return self.visit(state, node.circuit)
+
+    def visit_BitFlipChannel(
+        self, state: lowering.State[cirq.Circuit], node: cirq.BitFlipChannel
+    ):
+        p = node.gate.p
+        p_x = state.current_frame.push(py.Constant(p)).result
+        p_y = p_z = state.current_frame.push(py.Constant(0)).result
+        qubits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(
+            noise.stmts.SingleQubitPauliChannel(px=p_x, py=p_y, pz=p_z, qubits=qubits)
+        )
+
+    def visit_DepolarizingChannel(
+        self, state: lowering.State[cirq.Circuit], node: cirq.DepolarizingChannel
+    ):
+        p = state.current_frame.push(py.Constant(node.gate.p)).result
+        qubits = self.lower_qubit_getindices(state, node.qubits)
+        return state.current_frame.push(noise.stmts.Depolarize(p, qubits=qubits))
+
+    def visit_AsymmetricDepolarizingChannel(
+        self,
+        state: lowering.State[cirq.Circuit],
+        node: cirq.AsymmetricDepolarizingChannel,
+    ):
+        nqubits = node.gate.num_qubits()
+        if nqubits > 2:
+            raise lowering.BuildError(
+                "AsymmetricDepolarizingChannel applied to more than 2 qubits is not supported!"
+            )
+
+        if nqubits == 1:
+            qubits = self.lower_qubit_getindices(state, node.qubits)
+            p_x = state.current_frame.push(py.Constant(node.gate.p_x)).result
+            p_y = state.current_frame.push(py.Constant(node.gate.p_y)).result
+            p_z = state.current_frame.push(py.Constant(node.gate.p_z)).result
+            return state.current_frame.push(
+                noise.stmts.SingleQubitPauliChannel(p_x, p_y, p_z, qubits)
+            )
+
+        # NOTE: nqubits == 2
+        error_probs = node.gate.error_probabilities
+        probability_values = []
+        p0 = None
+        for key in self.two_qubit_paulis:
+            p = error_probs.get(key)
+
+            if p is None:
+                if p0 is None:
+                    p0 = state.current_frame.push(py.Constant(0)).result
+                p_ssa = p0
+            else:
+                p_ssa = state.current_frame.push(py.Constant(p)).result
+            probability_values.append(p_ssa)
+
+        probabilities = state.current_frame.push(
+            ilist.New(values=probability_values)
+        ).result
+
+        control, target = node.qubits
+        control_qarg = self.lower_qubit_getindices(state, (control,))
+        target_qarg = self.lower_qubit_getindices(state, (target,))
+
+        return state.current_frame.push(
+            noise.stmts.TwoQubitPauliChannel(
+                probabilities, controls=control_qarg, targets=target_qarg
+            )
+        )

bloqade/native/__init__.py

@@ -14,7 +14,6 @@ from .stdlib.simple import (
     ry as ry,
     rz as rz,
     u3 as u3,
-    rot as rot,
     s_dag as s_dag,
     shift as shift,
     sqrt_x as sqrt_x,

bloqade/native/_prelude.py

@@ -5,12 +5,12 @@ from kirin.passes import Default
 from kirin.prelude import structural_no_opt
 from typing_extensions import Doc
 
-from bloqade.squin import qubit
+from bloqade import qubit
 
-from .dialects import gates
+from .dialects import gate
 
 
-@ir.dialect_group(structural_no_opt.union([gates, qubit]))
+@ir.dialect_group(structural_no_opt.union([gate, qubit]))
 def kernel(self):
     """Compile a function to a native kernel."""
 

bloqade/native/dialects/gate/__init__.py

@@ -0,0 +1,2 @@
+from . import stmts as stmts
+from ._dialect import dialect as dialect

bloqade/native/dialects/gate/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("native.gate")