bloqade-circuit 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

bloqade/qasm2/dialects/expr/_emit.py

@@ -1,7 +1,6 @@
 from typing import Literal
 
 from kirin import interp
-from kirin.emit.exceptions import EmitError
 
 from bloqade.qasm2.parse import ast
 from bloqade.qasm2.types import QubitType
@@ -19,16 +18,18 @@ class EmitExpr(interp.MethodTable):
         self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: stmts.GateFunction
     ):
 
+        args: list[ast.Node] = []
         cparams, qparams = [], []
-        for arg in stmt.body.blocks[0].args[1:]:
-            name = frame.get(arg)
-            if not isinstance(name, ast.Name):
-                raise EmitError("expected ast.Name")
+        for arg in stmt.body.blocks[0].args:
+            assert arg.name is not None
+
+            args.append(ast.Name(id=arg.name))
             if arg.type.is_subseteq(QubitType):
-                qparams.append(name.id)
+                qparams.append(arg.name)
             else:
-                cparams.append(name.id)
-        emit.run_ssacfg_region(frame, stmt.body)
+                cparams.append(arg.name)
+
+        emit.run_ssacfg_region(frame, stmt.body, tuple(args))
         emit.output = ast.Gate(
             name=stmt.sym_name,
             cparams=cparams,

bloqade/qasm2/emit/base.py

@@ -36,8 +36,10 @@ class EmitQASM2Base(
         )
         return self
 
-    def new_frame(self, code: ir.Statement) -> EmitQASM2Frame:
-        return EmitQASM2Frame.from_func_like(code)
+    def initialize_frame(
+        self, code: ir.Statement, *, has_parent_access: bool = False
+    ) -> EmitQASM2Frame[StmtType]:
+        return EmitQASM2Frame(code, has_parent_access=has_parent_access)
 
     def run_method(
         self, method: ir.Method, args: tuple[ast.Node | None, ...]

bloqade/qasm2/emit/gate.py

@@ -86,17 +86,3 @@ class Func(interp.MethodTable):
     @interp.impl(func.ConstantNone)
     def ignore(self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt):
         return ()
-
-    @interp.impl(func.Function)
-    def emit_func(
-        self, emit: EmitQASM2Gate, frame: EmitQASM2Frame, stmt: func.Function
-    ):
-        emit.run_ssacfg_region(frame, stmt.body)
-        cparams, qparams = [], []
-        for arg in stmt.args:
-            if arg.type.is_subseteq(QubitType):
-                qparams.append(frame.get(arg))
-            else:
-                cparams.append(frame.get(arg))
-        emit.output = ast.Gate(stmt.sym_name, cparams, qparams, frame.body)
-        return ()

bloqade/qasm2/emit/main.py

@@ -24,7 +24,7 @@ class Func(interp.MethodTable):
     ):
         from bloqade.qasm2.dialects import glob, noise, parallel
 
-        emit.run_ssacfg_region(frame, stmt.body)
+        emit.run_ssacfg_region(frame, stmt.body, ())
         if emit.dialects.data.intersection(
             (parallel.dialect, glob.dialect, noise.dialect)
         ):
@@ -51,12 +51,14 @@ class Cf(interp.MethodTable):
         self, emit: EmitQASM2Main, frame: EmitQASM2Frame, stmt: cf.ConditionalBranch
     ):
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
-        body_frame = emit.new_frame(stmt)
-        body_frame.entries.update(frame.entries)
-        body_frame.set_values(
-            stmt.then_successor.args, frame.get_values(stmt.then_arguments)
-        )
-        emit.emit_block(body_frame, stmt.then_successor)
+
+        with emit.new_frame(stmt) as body_frame:
+            body_frame.entries.update(frame.entries)
+            body_frame.set_values(
+                stmt.then_successor.args, frame.get_values(stmt.then_arguments)
+            )
+            emit.emit_block(body_frame, stmt.then_successor)
+
         frame.body.append(
             ast.IfStmt(
                 cond,
@@ -91,15 +93,17 @@ class Scf(interp.MethodTable):
             )
 
         cond = emit.assert_node(ast.Cmp, frame.get(stmt.cond))
-        then_frame = emit.new_frame(stmt)
-        then_frame.entries.update(frame.entries)
-        emit.emit_block(then_frame, stmt.then_body.blocks[0])
-        frame.body.append(
-            ast.IfStmt(
-                cond,
-                body=then_frame.body,  # type: ignore
+
+        with emit.new_frame(stmt) as then_frame:
+            then_frame.entries.update(frame.entries)
+            emit.emit_block(then_frame, stmt.then_body.blocks[0])
+            frame.body.append(
+                ast.IfStmt(
+                    cond,
+                    body=then_frame.body,  # type: ignore
+                )
             )
-        )
+
         term = stmt.then_body.blocks[0].last_stmt
         if isinstance(term, scf.Yield):
             return then_frame.get_values(term.values)

bloqade/qasm2/emit/target.py

@@ -101,9 +101,7 @@ class QASM2:
 
         Py2QASM(entry.dialects)(entry)
         target_main = EmitQASM2Main(self.main_target)
-        target_main.run(
-            entry, tuple(ast.Name(name) for name in entry.arg_names[1:])
-        ).expect()
+        target_main.run(entry, ())
 
         main_program = target_main.output
         assert main_program is not None, f"failed to emit {entry.sym_name}"
@@ -133,9 +131,7 @@ class QASM2:
 
                 Py2QASM(fn.dialects)(fn)
 
-                target_gate.run(
-                    fn, tuple(ast.Name(name) for name in fn.arg_names[1:])
-                ).expect()
+                target_gate.run(fn, tuple(ast.Name(name) for name in fn.arg_names[1:]))
                 assert target_gate.output is not None, f"failed to emit {fn.sym_name}"
                 extra.append(target_gate.output)
 

bloqade/qasm2/glob.py

@@ -11,7 +11,7 @@ from .dialects import glob
 
 @wraps(glob.UGate)
 def u(
-    theta: float, phi: float, lam: float, registers: ilist.IList[QReg, Any] | list
+    registers: ilist.IList[QReg, Any] | list, theta: float, phi: float, lam: float
 ) -> None:
     """Apply a U gate to all qubits in the input registers.
 

bloqade/qasm2/parse/lowering.py

@@ -1,3 +1,5 @@
+import os
+import pathlib
 from typing import Any
 from dataclasses import field, dataclass
 
@@ -17,6 +19,119 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     hint_show_lineno: bool = field(default=True, kw_only=True)
     stacktrace: bool = field(default=True, kw_only=True)
 
+    def loads(
+        self,
+        source: str,
+        kernel_name: str,
+        *,
+        returns: str | None = None,
+        globals: dict[str, Any] | None = None,
+        file: str | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ) -> ir.Method:
+        from ..parse import loads
+
+        # TODO: add source info
+        stmt = loads(source)
+
+        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:
+            try:
+                self.visit(state, stmt)
+                # append return statement with the return values
+                if returns is not None:
+                    return_value = frame.get(returns)
+                    if return_value is None:
+                        raise lowering.BuildError(f"Cannot find return value {returns}")
+                else:
+                    return_value = func.ConstantNone()
+
+                return_node = frame.push(func.Return(value_or_stmt=return_value))
+
+            except lowering.BuildError as e:
+                hint = state.error_hint(
+                    e,
+                    max_lines=self.max_lines,
+                    indent=self.hint_indent,
+                    show_lineno=self.hint_show_lineno,
+                )
+                if self.stacktrace:
+                    raise Exception(
+                        f"{e.args[0]}\n\n{hint}",
+                        *e.args[1:],
+                    ) from e
+                else:
+                    e.args = (hint,)
+                    raise e
+
+            region = frame.curr_region
+
+        if compactify:
+            from kirin.rewrite import Walk, CFGCompactify
+
+            Walk(CFGCompactify()).rewrite(region)
+
+        code = func.Function(
+            sym_name=kernel_name,
+            signature=func.Signature((), return_node.value.type),
+            body=region,
+        )
+
+        return ir.Method(
+            mod=None,
+            py_func=None,
+            sym_name=kernel_name,
+            arg_names=[],
+            dialects=self.dialects,
+            code=code,
+        )
+
+    def loadfile(
+        self,
+        file: str | pathlib.Path,
+        *,
+        kernel_name: str | None = None,
+        returns: str | None = None,
+        globals: dict[str, Any] | None = None,
+        lineno_offset: int = 0,
+        col_offset: int = 0,
+        compactify: bool = True,
+    ) -> ir.Method:
+        if isinstance(file, str):
+            file = pathlib.Path(*os.path.split(file))
+
+        if not file.is_file() or not file.name.endswith(".qasm"):
+            raise ValueError("File must be a .qasm file")
+
+        kernel_name = (
+            file.name.replace(".qasm", "") if kernel_name is None else kernel_name
+        )
+
+        with file.open("r") as f:
+            source = f.read()
+
+        return self.loads(
+            source,
+            kernel_name,
+            returns=returns,
+            globals=globals,
+            file=str(file),
+            lineno_offset=lineno_offset,
+            col_offset=col_offset,
+            compactify=compactify,
+        )
+
     def run(
         self,
         stmt: ast.Node,
@@ -85,6 +200,10 @@ class QASM2(lowering.LoweringABC[ast.Node]):
             stmt = expr.ConstInt(value=value)
         elif isinstance(value, float):
             stmt = expr.ConstFloat(value=value)
+        else:
+            raise lowering.BuildError(
+                f"Expected value of type float or int, got {type(value)}."
+            )
         state.current_frame.push(stmt)
         return stmt.result
 
@@ -99,6 +218,8 @@ class QASM2(lowering.LoweringABC[ast.Node]):
             dialects = ["qasm2.core", "qasm2.uop", "qasm2.expr"]
         elif isinstance(node.header, ast.Kirin):
             dialects = node.header.dialects
+        else:
+            raise lowering.BuildError(f"Unexpected node header {node.header}")
 
         for dialect in dialects:
             if dialect not in allowed:
@@ -278,7 +399,7 @@ class QASM2(lowering.LoweringABC[ast.Node]):
     def visit_UnaryOp(self, state: lowering.State[ast.Node], node: ast.UnaryOp):
         if node.op == "-":
             stmt = expr.Neg(value=state.lower(node.operand).expect_one())
-            return stmt.result
+            return state.current_frame.push(stmt).result
         else:
             return state.lower(node.operand).expect_one()
 
@@ -295,6 +416,8 @@ class QASM2(lowering.LoweringABC[ast.Node]):
             stmt = core.QRegGet(reg, addr.result)
         elif reg.type.is_subseteq(CRegType):
             stmt = core.CRegGet(reg, addr.result)
+        else:
+            raise lowering.BuildError(f"Unexpected register type {reg.type}")
         return state.current_frame.push(stmt).result
 
     def visit_Call(self, state: lowering.State[ast.Node], node: ast.Call):

bloqade/qasm2/passes/glob.py

@@ -4,7 +4,7 @@ which converts global gates to single qubit gates.
 """
 
 from kirin import ir
-from kirin.rewrite import cse, dce, walk, result
+from kirin.rewrite import abc, cse, dce, walk
 from kirin.passes.abc import Pass
 from kirin.passes.fold import Fold
 from kirin.rewrite.fixpoint import Fixpoint
@@ -54,7 +54,7 @@ class GlobalToUOP(Pass):
         frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
         return GlobalToUOpRule(frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 
@@ -106,7 +106,7 @@ class GlobalToParallel(Pass):
         frame, _ = address.AddressAnalysis(mt.dialects).run_analysis(mt)
         return GlobalToParallelRule(frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         rewriter = walk.Walk(self.generate_rule(mt))
         result = rewriter.rewrite(mt.code)
 

bloqade/qasm2/passes/lift_qubits.py

@@ -0,0 +1,26 @@
+from kirin import ir
+from kirin.passes import Pass
+from kirin.rewrite import (
+    Walk,
+    Chain,
+    Fixpoint,
+    ConstantFold,
+    CommonSubexpressionElimination,
+)
+from kirin.passes.hint_const import HintConst
+
+from bloqade.qasm2.rewrite.insert_qubits import InsertGetQubit
+
+
+class LiftQubits(Pass):
+    """This pass lifts the creation of qubits to the block where the register is defined."""
+
+    def unsafe_run(self, mt: ir.Method):
+        result = Walk(InsertGetQubit()).rewrite(mt.code)
+        result = HintConst(self.dialects).unsafe_run(mt).join(result)
+        result = (
+            Fixpoint(Walk(Chain(ConstantFold(), CommonSubexpressionElimination())))
+            .rewrite(mt.code)
+            .join(result)
+        )
+        return result

bloqade/qasm2/passes/noise.py

@@ -4,16 +4,13 @@ from kirin import ir
 from kirin.passes import Pass
 from kirin.rewrite import (
     Walk,
-    Chain,
     Fixpoint,
-    ConstantFold,
     DeadCodeElimination,
-    CommonSubexpressionElimination,
 )
-from kirin.rewrite.result import RewriteResult
 
 from bloqade.noise import native
 from bloqade.analysis import address
+from bloqade.qasm2.passes.lift_qubits import LiftQubits
 from bloqade.qasm2.rewrite.heuristic_noise import NoiseRewriteRule
 
 
@@ -38,24 +35,19 @@ class NoisePass(Pass):
         self.address_analysis = address.AddressAnalysis(self.dialects)
 
     def unsafe_run(self, mt: ir.Method):
-        result = RewriteResult()
-
-        frame, res = self.address_analysis.run_analysis(mt, no_raise=False)
+        result = LiftQubits(self.dialects).unsafe_run(mt)
+        frame, _ = self.address_analysis.run_analysis(mt, no_raise=self.no_raise)
         result = (
             Walk(
                 NoiseRewriteRule(
                     address_analysis=frame.entries,
                     noise_model=self.noise_model,
                     gate_noise_params=self.gate_noise_params,
-                )
+                ),
+                reverse=True,
             )
             .rewrite(mt.code)
             .join(result)
         )
-        rule = Chain(
-            ConstantFold(),
-            DeadCodeElimination(),
-            CommonSubexpressionElimination(),
-        )
-        result = Fixpoint(Walk(rule)).rewrite(mt.code).join(result)
+        result = Fixpoint(Walk(DeadCodeElimination())).rewrite(mt.code).join(result)
         return result

bloqade/qasm2/passes/parallel.py

@@ -16,7 +16,7 @@ from kirin.rewrite import (
     ConstantFold,
     DeadCodeElimination,
     CommonSubexpressionElimination,
-    result,
+    abc,
 )
 from kirin.analysis import const
 
@@ -84,7 +84,7 @@ class ParallelToUOp(Pass):
 
         return ParallelToUOpRule(id_map=id_map, address_analysis=frame.entries)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         result = Walk(self.generate_rule(mt)).rewrite(mt.code)
         rule = Chain(
             ConstantFold(),
@@ -140,7 +140,7 @@ class UOpToParallel(Pass):
     def __post_init__(self):
         self.constprop = const.Propagate(self.dialects)
 
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
         result = Walk(RaiseRegisterRule()).rewrite(mt.code)
 
         # do not run the parallelization because registers are not at the top

bloqade/qasm2/passes/py2qasm.py

@@ -4,8 +4,7 @@ from kirin import ir
 from kirin.passes import Pass
 from kirin.rewrite import Walk, Fixpoint
 from kirin.dialects import py, math
-from kirin.rewrite.abc import RewriteRule
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.qasm2.dialects import core, expr
 

bloqade/qasm2/passes/qasm2py.py

@@ -6,8 +6,7 @@ from kirin import ir
 from kirin.passes import Pass
 from kirin.rewrite import Walk, Fixpoint
 from kirin.dialects import py, math
-from kirin.rewrite.abc import RewriteRule
-from kirin.rewrite.result import RewriteResult
+from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.qasm2.dialects import core, expr
 

bloqade/qasm2/rewrite/desugar.py

@@ -2,27 +2,27 @@ from dataclasses import dataclass
 
 from kirin import ir
 from kirin.passes import Pass
-from kirin.rewrite import abc, walk, result
+from kirin.rewrite import abc, walk
 from kirin.dialects import py
 
 from bloqade.qasm2.dialects import core
 
 
 class IndexingDesugarRule(abc.RewriteRule):
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if isinstance(node, py.indexing.GetItem):
             if node.obj.type.is_subseteq(core.QRegType):
                 node.replace_by(core.QRegGet(reg=node.obj, idx=node.index))
-                return result.RewriteResult(has_done_something=True)
+                return abc.RewriteResult(has_done_something=True)
             elif node.obj.type.is_subseteq(core.CRegType):
                 node.replace_by(core.CRegGet(reg=node.obj, idx=node.index))
-                return result.RewriteResult(has_done_something=True)
+                return abc.RewriteResult(has_done_something=True)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
 
 @dataclass
 class IndexingDesugarPass(Pass):
-    def unsafe_run(self, mt: ir.Method) -> result.RewriteResult:
+    def unsafe_run(self, mt: ir.Method) -> abc.RewriteResult:
 
         return walk.Walk(IndexingDesugarRule()).rewrite(mt.code)

bloqade/qasm2/rewrite/glob.py

@@ -2,7 +2,7 @@ from typing import Dict, List
 from dataclasses import dataclass
 
 from kirin import ir
-from kirin.rewrite import abc, result
+from kirin.rewrite import abc
 from kirin.dialects import py, ilist
 
 from bloqade import qasm2
@@ -47,18 +47,18 @@ class GlobalRewriteBase:
 @dataclass
 class GlobalToParallelRule(abc.RewriteRule, GlobalRewriteBase):
 
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if type(node) in glob.dialect.stmts:
             return getattr(self, f"rewrite_{node.name}")(node)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
     def rewrite_ugate(self, node: glob.UGate):
 
         new_stmts, qubit_ssa = self.get_qubit_ssa(node)
 
         if qubit_ssa is None:
-            return result.RewriteResult()
+            return abc.RewriteResult()
 
         new_stmts.append(qargs := ilist.New(values=qubit_ssa))
         new_stmts.append(
@@ -72,24 +72,24 @@ class GlobalToParallelRule(abc.RewriteRule, GlobalRewriteBase):
 
         node.delete()
 
-        return result.RewriteResult(has_done_something=True)
+        return abc.RewriteResult(has_done_something=True)
 
 
 @dataclass
 class GlobalToUOpRule(abc.RewriteRule, GlobalRewriteBase):
 
-    def rewrite_Statement(self, node: ir.Statement) -> result.RewriteResult:
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         if type(node) in glob.dialect.stmts:
             return getattr(self, f"rewrite_{node.name}")(node)
 
-        return result.RewriteResult()
+        return abc.RewriteResult()
 
     def rewrite_ugate(self, node: glob.UGate):
 
         new_stmts, qubit_ssa = self.get_qubit_ssa(node)
 
         if qubit_ssa is None:
-            return result.RewriteResult()
+            return abc.RewriteResult()
 
         for qarg in qubit_ssa:
             new_stmts.append(
@@ -100,4 +100,4 @@ class GlobalToUOpRule(abc.RewriteRule, GlobalRewriteBase):
             stmt.insert_before(node)
 
         node.delete()
-        return result.RewriteResult(has_done_something=True)
+        return abc.RewriteResult(has_done_something=True)