Clean up import path (#279)

This PR address #276, remove all `from xxx import *`

General pattern:

1. to access the Statement definition: 

```python 
from stim.dialects import  collapse, auxiliary, gate

gate.CX # the statement def 
auxiliary.ObservableInclude # the statement def
```

2. for general user access, they should not touch statement, so:
```python
from bloqade import stim

stim.cx # wrapper 
stim.observable_include # wrapper
```

Author: kaihsin

src/bloqade/qasm2/dialects/core/__init__.py

@@ -1,3 +1,11 @@
 from . import _emit as _emit, address as address, _typeinfer as _typeinfer
-from .stmts import *  # noqa: F403
+from .stmts import (
+    Reset as Reset,
+    CRegEq as CRegEq,
+    CRegGet as CRegGet,
+    CRegNew as CRegNew,
+    Measure as Measure,
+    QRegGet as QRegGet,
+    QRegNew as QRegNew,
+)
 from ._dialect import dialect as dialect

src/bloqade/qasm2/dialects/expr/__init__.py

@@ -1,3 +1,20 @@
 from . import _emit as _emit, _interp as _interp, _from_python as _from_python
-from .stmts import *  # noqa: F403
+from .stmts import (
+    Add as Add,
+    Cos as Cos,
+    Div as Div,
+    Exp as Exp,
+    Log as Log,
+    Mul as Mul,
+    Neg as Neg,
+    Pow as Pow,
+    Sin as Sin,
+    Sub as Sub,
+    Tan as Tan,
+    Sqrt as Sqrt,
+    ConstPI as ConstPI,
+    ConstInt as ConstInt,
+    ConstFloat as ConstFloat,
+    GateFunction as GateFunction,
+)
 from ._dialect import dialect as dialect

src/bloqade/qasm2/dialects/uop/__init__.py

@@ -1,4 +1,40 @@
-from . import _emit as _emit, stmts as stmts
-from .stmts import *  # noqa: F403
+from . import _emit as _emit, stmts as stmts, schedule as schedule
+from .stmts import (
+    CH as CH,
+    CU as CU,
+    CX as CX,
+    CY as CY,
+    CZ as CZ,
+    RX as RX,
+    RY as RY,
+    RZ as RZ,
+    SX as SX,
+    U1 as U1,
+    U2 as U2,
+    CCX as CCX,
+    CRX as CRX,
+    CRY as CRY,
+    CRZ as CRZ,
+    CSX as CSX,
+    CU1 as CU1,
+    CU3 as CU3,
+    RXX as RXX,
+    RZZ as RZZ,
+    H as H,
+    S as S,
+    T as T,
+    X as X,
+    Y as Y,
+    Z as Z,
+    Id as Id,
+    Sdag as Sdag,
+    Swap as Swap,
+    Tdag as Tdag,
+    CSwap as CSwap,
+    SXdag as SXdag,
+    UGate as UGate,
+    Barrier as Barrier,
+    SingleQubitGate as SingleQubitGate,
+    TwoQubitCtrlGate as TwoQubitCtrlGate,
+)
 from ._dialect import dialect as dialect
-from .schedule import *  # noqa: F403

src/bloqade/qasm2/dialects/uop/schedule.py

@@ -8,7 +8,7 @@
 
 
 @dialect.register(key="qasm2.schedule.dag")
-class UOp(interp.MethodTable):
+class UOpSchedule(interp.MethodTable):
 
     @interp.impl(stmts.Id)
     @interp.impl(stmts.SXdag)

src/bloqade/qasm2/rewrite/desugar.py

@@ -5,16 +5,17 @@
 from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
+from bloqade.qasm2 import types
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
     def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
-            if node.obj.type.is_subseteq(core.QRegType):
+            if node.obj.type.is_subseteq(types.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
-            elif node.obj.type.is_subseteq(core.CRegType):
+            elif node.obj.type.is_subseteq(types.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
                 return abc.RewriteResult(has_done_something=True)
 

src/bloqade/squin/noise/__init__.py

@@ -1,27 +1,8 @@
-# Put all the proper wrappers here
-
-from kirin.lowering import wraps as _wraps
-
-from bloqade.squin.op.types import Op
-
 from . import stmts as stmts
-
-
-@_wraps(stmts.PauliError)
-def pauli_error(basis: Op, p: float) -> Op: ...
-
-
-@_wraps(stmts.PPError)
-def pp_error(op: Op, p: float) -> Op: ...
-
-
-@_wraps(stmts.Depolarize)
-def depolarize(n_qubits: int, p: float) -> Op: ...
-
-
-@_wraps(stmts.PauliChannel)
-def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
-
-
-@_wraps(stmts.QubitLoss)
-def qubit_loss(p: float) -> Op: ...
+from ._dialect import dialect as dialect
+from ._wrapper import (
+    pp_error as pp_error,
+    depolarize as depolarize,
+    qubit_loss as qubit_loss,
+    pauli_channel as pauli_channel,
+)

src/bloqade/squin/noise/_wrapper.py

@@ -0,0 +1,25 @@
+from kirin.lowering import wraps
+
+from bloqade.squin.op.types import Op
+
+from . import stmts
+
+
+@wraps(stmts.PauliError)
+def pauli_error(basis: Op, p: float) -> Op: ...
+
+
+@wraps(stmts.PPError)
+def pp_error(op: Op, p: float) -> Op: ...
+
+
+@wraps(stmts.Depolarize)
+def depolarize(n_qubits: int, p: float) -> Op: ...
+
+
+@wraps(stmts.PauliChannel)
+def pauli_channel(n_qubits: int, params: tuple[float, ...]) -> Op: ...
+
+
+@wraps(stmts.QubitLoss)
+def qubit_loss(p: float) -> Op: ...

src/bloqade/squin/op/__init__.py

@@ -1,162 +1,36 @@
-from kirin import ir as _ir
-from kirin.prelude import structural_no_opt as _structural_no_opt
-from kirin.lowering import wraps as _wraps
-
 from . import stmts as stmts, types as types, rewrite as rewrite
+from .stdlib import (
+    ch as ch,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    cphase as cphase,
+)
 from .traits import Unitary as Unitary, MaybeUnitary as MaybeUnitary
 from ._dialect import dialect as dialect
-
-
-@_wraps(stmts.Kron)
-def kron(lhs: types.Op, rhs: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Mult)
-def mult(lhs: types.Op, rhs: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Scale)
-def scale(op: types.Op, factor: complex) -> types.Op: ...
-
-
-@_wraps(stmts.Adjoint)
-def adjoint(op: types.Op) -> types.Op: ...
-
-
-@_wraps(stmts.Control)
-def control(op: types.Op, *, n_controls: int) -> types.Op:
-    """
-    Create a controlled operator.
-
-    Note, that when considering atom loss, the operator will not be applied if
-    any of the controls has been lost.
-
-    Args:
-        operator: The operator to apply under the control.
-        n_controls: The number qubits to be used as control.
-
-    Returns:
-        Operator
-    """
-    ...
-
-
-@_wraps(stmts.Identity)
-def identity(*, sites: int) -> types.Op: ...
-
-
-@_wraps(stmts.Rot)
-def rot(axis: types.Op, angle: float) -> types.Op: ...
-
-
-@_wraps(stmts.ShiftOp)
-def shift(theta: float) -> types.Op: ...
-
-
-@_wraps(stmts.PhaseOp)
-def phase(theta: float) -> types.Op: ...
-
-
-@_wraps(stmts.X)
-def x() -> types.Op: ...
-
-
-@_wraps(stmts.Y)
-def y() -> types.Op: ...
-
-
-@_wraps(stmts.Z)
-def z() -> types.Op: ...
-
-
-@_wraps(stmts.H)
-def h() -> types.Op: ...
-
-
-@_wraps(stmts.S)
-def s() -> types.Op: ...
-
-
-@_wraps(stmts.T)
-def t() -> types.Op: ...
-
-
-@_wraps(stmts.P0)
-def p0() -> types.Op: ...
-
-
-@_wraps(stmts.P1)
-def p1() -> types.Op: ...
-
-
-@_wraps(stmts.Sn)
-def spin_n() -> types.Op: ...
-
-
-@_wraps(stmts.Sp)
-def spin_p() -> types.Op: ...
-
-
-@_wraps(stmts.U3)
-def u(theta: float, phi: float, lam: float) -> types.Op: ...
-
-
-@_wraps(stmts.PauliString)
-def pauli_string(*, string: str) -> types.Op: ...
-
-
-# stdlibs
-@_ir.dialect_group(_structural_no_opt.add(dialect))
-def op(self):
-    def run_pass(method):
-        pass
-
-    return run_pass
-
-
-@op
-def rx(theta: float) -> types.Op:
-    """Rotation X gate."""
-    return rot(x(), theta)
-
-
-@op
-def ry(theta: float) -> types.Op:
-    """Rotation Y gate."""
-    return rot(y(), theta)
-
-
-@op
-def rz(theta: float) -> types.Op:
-    """Rotation Z gate."""
-    return rot(z(), theta)
-
-
-@op
-def cx() -> types.Op:
-    """Controlled X gate."""
-    return control(x(), n_controls=1)
-
-
-@op
-def cy() -> types.Op:
-    """Controlled Y gate."""
-    return control(y(), n_controls=1)
-
-
-@op
-def cz() -> types.Op:
-    """Control Z gate."""
-    return control(z(), n_controls=1)
-
-
-@op
-def ch() -> types.Op:
-    """Control H gate."""
-    return control(h(), n_controls=1)
-
-
-@op
-def cphase(theta: float) -> types.Op:
-    """Control Phase gate."""
-    return control(phase(theta), n_controls=1)
+from ._wrapper import (
+    h as h,
+    s as s,
+    t as t,
+    u as u,
+    x as x,
+    y as y,
+    z as z,
+    p0 as p0,
+    p1 as p1,
+    rot as rot,
+    kron as kron,
+    mult as mult,
+    phase as phase,
+    scale as scale,
+    shift as shift,
+    spin_n as spin_n,
+    spin_p as spin_p,
+    adjoint as adjoint,
+    control as control,
+    identity as identity,
+    pauli_string as pauli_string,
+)