Adding tweezer task. (#16)

* updating interface

* adding task function for tweezer calls inside Auto retion

* adding new statement representing an unassigned tweezer function for inside Auto

* allow for  as well as invoke

Author: weinbe58

.gitignore

@@ -1,3 +1,7 @@
+# misc
+
+debug/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

src/bloqade/shuttle/dialects/schedule/__init__.py

@@ -6,6 +6,7 @@
     Auto as Auto,
     ExecutableRegion as ExecutableRegion,
     NewDeviceFunction as NewDeviceFunction,
+    NewTweezerTask as NewTweezerTask,
     Parallel as Parallel,
     Reverse as Reverse,
 )

src/bloqade/shuttle/dialects/schedule/_interface.py

@@ -1,11 +1,13 @@
-from typing import Any, ParamSpec
+from typing import Any, ContextManager, ParamSpec
 
 from kirin import ir
 from kirin.dialects import ilist
 from kirin.lowering import wraps as _wraps
 
 from .stmts import (
+    Auto,
     NewDeviceFunction,
+    Parallel,
     Reverse,
 )
 from .types import DeviceFunction, ReverseDeviceFunction
@@ -43,3 +45,11 @@ def reverse(
 
     """
     ...
+
+
+@_wraps(Parallel)
+def parallel() -> ContextManager: ...
+
+
+@_wraps(Auto)
+def auto() -> ContextManager: ...

src/bloqade/shuttle/dialects/schedule/stmts.py

@@ -11,6 +11,14 @@
 NyTones = types.TypeVar("NyTones")
 
 
+@statement(dialect=dialect)
+class NewTweezerTask(ir.Statement):
+    name = "tweezer_task"
+    traits = frozenset({ir.Pure()})
+    move_fn: ir.SSAValue = info.argument(types.MethodType)
+    result: ir.ResultValue = info.result(DeviceFunctionType)
+
+
 @statement(dialect=dialect)
 class NewDeviceFunction(ir.Statement):
     name = "device_function"

src/bloqade/shuttle/passes/schedule2path.py

@@ -10,6 +10,7 @@
 
 from bloqade.shuttle.rewrite.schedule2path import (
     Canonicalize,
+    RewriteAutoInvoke,
     RewriteDeviceCall,
     RewriteScheduleRegion,
 )
@@ -20,7 +21,11 @@ class ScheduleToPath(Pass):
 
     def unsafe_run(self, mt: ir.Method):
         result = Fixpoint(Walk(Canonicalize())).rewrite(mt.code)
-        result = Walk(RewriteDeviceCall()).rewrite(mt.code).join(result)
+        result = (
+            Walk(Chain(RewriteAutoInvoke(), RewriteDeviceCall()))
+            .rewrite(mt.code)
+            .join(result)
+        )
         result = Walk(RewriteScheduleRegion()).rewrite(mt.code).join(result)
         result = (
             Fixpoint(

src/bloqade/shuttle/rewrite/auto_scheduler.py

@@ -1,5 +1,5 @@
 from kirin import ir
-from kirin.dialects import func
+from kirin.dialects import func, py
 from kirin.rewrite import abc
 
 from bloqade.shuttle.dialects import path, schedule
@@ -24,6 +24,31 @@ def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
         return abc.RewriteResult(has_done_something=True)
 
 
+class RewriteAutoInvoke(abc.RewriteRule):
+    def rewrite_Statement(self, node: ir.Statement) -> abc.RewriteResult:
+        if not isinstance(node.parent_stmt, schedule.Auto):
+            return abc.RewriteResult()
+
+        if isinstance(node, func.Invoke):
+            (callee_stmt := py.Constant(node.callee)).insert_before(node)
+            callee_ssa = callee_stmt.result
+        elif isinstance(node, func.Call):
+            callee_ssa = node.callee
+        else:
+            return abc.RewriteResult()
+
+        (tweezer_task := schedule.NewTweezerTask(move_fn=callee_ssa)).insert_before(
+            node
+        )
+        (path.Gen(tweezer_task.result, node.inputs, kwargs=node.kwargs)).insert_before(
+            node
+        )
+
+        node.delete()
+
+        return abc.RewriteResult(has_done_something=True)
+
+
 class RewriteScheduleRegion(abc.RewriteRule):
     CLASSES = {
         schedule.Auto: path.Auto,