FromPythonWith trait (#142)

Author: Roger-luo

src/kirin/ir/__init__.py

@@ -33,9 +33,11 @@
     NoTerminator as NoTerminator,
     SSACFGRegion as SSACFGRegion,
     FromPythonCall as FromPythonCall,
+    FromPythonWith as FromPythonWith,
     IsolatedFromAbove as IsolatedFromAbove,
     SymbolOpInterface as SymbolOpInterface,
     PythonLoweringTrait as PythonLoweringTrait,
     CallableStmtInterface as CallableStmtInterface,
+    FromPythonWithSingleItem as FromPythonWithSingleItem,
 )
 from kirin.ir.dialect import Dialect as Dialect

src/kirin/ir/traits/__init__.py

@@ -18,3 +18,7 @@
 )
 from .lowering.call import FromPythonCall as FromPythonCall
 from .region.ssacfg import SSACFGRegion as SSACFGRegion
+from .lowering.context import (
+    FromPythonWith as FromPythonWith,
+    FromPythonWithSingleItem as FromPythonWithSingleItem,
+)

src/kirin/ir/traits/lowering/call.py

@@ -2,8 +2,6 @@
 from typing import TYPE_CHECKING, TypeVar
 from dataclasses import dataclass
 
-from kirin.exceptions import DialectLoweringError
-
 from ..abc import PythonLoweringTrait
 
 if TYPE_CHECKING:
@@ -16,60 +14,24 @@
 
 @dataclass(frozen=True)
 class FromPythonCall(PythonLoweringTrait[StatementType, ast.Call]):
+    """Trait for customizing lowering of Python calls to a statement.
+
+    Declared in a statement definition to indicate that the statement can be
+    constructed from a Python call (i.e., a function call `ast.Call` in the
+    Python AST).
+
+    Subclassing this trait allows for customizing the lowering of Python calls
+    to the statement. The `lower` method should be implemented to parse the
+    arguments from the Python call and construct the statement instance.
+    """
 
     def lower(
         self, stmt: type[StatementType], state: "LoweringState", node: ast.Call
     ) -> "Result":
-        from kirin.decl import fields
-        from kirin.lowering import Result
-        from kirin.dialects.py.data import PyAttr
-
-        fs = fields(stmt)
-        stmt_std_arg_names = fs.std_args.keys()
-        stmt_kw_args_name = fs.kw_args.keys()
-        stmt_attr_prop_names = fs.attr_or_props
-        stmt_required_names = fs.required_names
-        stmt_group_arg_names = fs.group_arg_names
-        args, kwargs = {}, {}
-        for name, value in zip(stmt_std_arg_names, node.args):
-            self._parse_arg(stmt_group_arg_names, state, args, name, value)
-        for kw in node.keywords:
-            if not isinstance(kw.arg, str):
-                raise DialectLoweringError("Expected string for keyword argument name")
-
-            arg: str = kw.arg
-            if arg in node.args:
-                raise DialectLoweringError(
-                    f"Keyword argument {arg} is already present in positional arguments"
-                )
-            elif arg in stmt_std_arg_names or arg in stmt_kw_args_name:
-                self._parse_arg(stmt_group_arg_names, state, kwargs, kw.arg, kw.value)
-            elif arg in stmt_attr_prop_names:
-                if not isinstance(kw.value, ast.Constant):
-                    raise DialectLoweringError(
-                        f"Expected constant for attribute or property {arg}"
-                    )
-                kwargs[arg] = PyAttr(kw.value.value)
-            else:
-                raise DialectLoweringError(f"Unexpected keyword argument {arg}")
-
-        for name in stmt_required_names:
-            if name not in args and name not in kwargs:
-                raise DialectLoweringError(f"Missing required argument {name}")
-
-        return Result(state.append_stmt(stmt(*args.values(), **kwargs)))
+        return state.default_Call_lower(stmt, node)
 
-    @staticmethod
-    def _parse_arg(
-        group_names: set[str],
-        state: "LoweringState",
-        target: dict,
-        name: str,
-        value: ast.AST,
-    ):
-        if name in group_names:
-            if not isinstance(value, ast.Tuple):
-                raise DialectLoweringError(f"Expected tuple for group argument {name}")
-            target[name] = tuple(state.visit(elem).expect_one() for elem in value.elts)
-        else:
-            target[name] = state.visit(value).expect_one()
+    def verify(self, stmt: "Statement"):
+        assert len(stmt.regions) == 0, "FromPythonCall statements cannot have regions"
+        assert (
+            len(stmt.successors) == 0
+        ), "FromPythonCall statements cannot have successors"

src/kirin/ir/traits/lowering/context.py

@@ -0,0 +1,105 @@
+"""Traits for customizing lowering of Python `with` syntax to a statement.
+"""
+
+import ast
+from typing import TYPE_CHECKING, TypeVar
+from dataclasses import dataclass
+
+from kirin.exceptions import DialectLoweringError
+
+from ..abc import PythonLoweringTrait
+
+if TYPE_CHECKING:
+    from kirin.ir import Statement
+    from kirin.lowering import Result, LoweringState
+
+StatementType = TypeVar("StatementType", bound="Statement")
+
+
+@dataclass(frozen=True)
+class FromPythonWith(PythonLoweringTrait[StatementType, ast.With]):
+    """Trait for customizing lowering of Python with statements to a statement.
+
+    Subclassing this trait allows for customizing the lowering of Python with
+    statements to the statement. The `lower` method should be implemented to parse
+    the arguments from the Python with statement and construct the statement instance.
+    """
+
+    pass
+
+
+@dataclass(frozen=True)
+class FromPythonWithSingleItem(FromPythonWith[StatementType]):
+    """Trait for customizing lowering of the following Python with syntax to a statement:
+
+    ```python
+    with <stmt>[ as <name>]:
+        <body>
+    ```
+
+    where `<stmt>` is the statement being lowered, `<name>` is an optional name for the result
+    of the statement, and `<body>` is the body of the with statement. The optional `as <name>`
+    is not valid when the statement has no results.
+
+    This syntax is slightly different from the standard Python `with` statement in that
+    `<name>` refers to the result of the statement, not the context manager. Thus typically
+    one sould access `<name>` in `<body>` to use the result of the statement.
+
+    In some cases, however, `<name>` may be used as a reference of a special value `self` that
+    is passed to the `<body>` of the statement. This is useful for statements that have a similar
+    behavior to a closure.
+    """
+
+    def lower(
+        self, stmt: type[StatementType], state: "LoweringState", node: ast.With
+    ) -> "Result":
+        from kirin import lowering
+        from kirin.decl import fields
+
+        fs = fields(stmt)
+        if len(fs.regions) != 1:
+            raise DialectLoweringError(
+                "Expected exactly one region in statement declaration"
+            )
+
+        if len(node.items) != 1:
+            raise DialectLoweringError("Expected exactly one item in statement")
+
+        item, body = node.items[0], node.body
+        if not isinstance(item.context_expr, ast.Call):
+            raise DialectLoweringError(
+                f"Expected context expression to be a call for with {stmt.name}"
+            )
+
+        body_frame = lowering.Frame.from_stmts(body, state, parent=state.current_frame)
+        state.push_frame(body_frame)
+        state.exhaust()
+        state.pop_frame()
+
+        args, kwargs = state.default_Call_inputs(stmt, item.context_expr)
+        (region_name,) = fs.regions
+        kwargs[region_name] = body_frame.current_region
+        results = state.append_stmt(stmt(*args.values(), **kwargs)).results
+        if len(results) == 0:
+            return lowering.Result()
+        elif len(results) > 1:
+            raise DialectLoweringError(
+                f"Expected exactly one result or no result from statement {stmt.name}"
+            )
+
+        result = results[0]
+        if item.optional_vars is not None and isinstance(item.optional_vars, ast.Name):
+            result.name = item.optional_vars.id
+            state.current_frame.defs[result.name] = result
+        return lowering.Result(result)
+
+    def verify(self, stmt: "Statement"):
+        assert (
+            len(stmt.regions) == 1
+        ), "FromPythonWithSingleItem statements must have one region"
+        assert (
+            len(stmt.successors) == 0
+        ), "FromPythonWithSingleItem statements cannot have successors"
+        assert (
+            len(stmt.results) <= 1
+        ), "FromPythonWithSingleItem statements can have at most one result"

src/kirin/lowering/state.py

@@ -112,11 +112,38 @@ def visit(self, node: ast.AST) -> Result:
             # it will be called first before __dispatch_Call
             # because "Call" exists in self.registry
             return self.__dispatch_Call(node)
+        elif isinstance(node, ast.With):
+            return self.__dispatch_With(node)
         return super().visit(node)
 
     def generic_visit(self, node: ast.AST):
         raise DialectLoweringError(f"unsupported ast node {type(node)}:")
 
+    def __dispatch_With(self, node: ast.With):
+        if len(node.items) != 1:
+            raise DialectLoweringError("expected exactly one item in with statement")
+
+        item = node.items[0]
+        if not isinstance(item.context_expr, ast.Call):
+            raise DialectLoweringError("expected context expression to be a call")
+
+        global_callee_result = self.get_global_nothrow(item.context_expr.func)
+        if global_callee_result is None:
+            raise DialectLoweringError("cannot find call func in with context")
+
+        global_callee = global_callee_result.unwrap()
+        if not issubclass(global_callee, Statement):
+            raise DialectLoweringError("expected callee to be a statement")
+
+        if (
+            trait := global_callee.get_trait(traits.FromPythonWithSingleItem)
+        ) is not None:
+            return trait.lower(global_callee, self, node)
+
+        raise DialectLoweringError(
+            "unsupported callee, missing FromPythonWithSingleItem trait"
+        )
+
     def __dispatch_Call(self, node: ast.Call):
         # 1. try to lookup global statement object
         # 2. lookup local values
@@ -196,6 +223,63 @@ def __lower_Call_local(self, node: ast.Call) -> Result:
             return self.registry["Call_local"].lower_Call_local(self, callee, node)
         raise DialectLoweringError("`lower_Call_local` not implemented")
 
+    def default_Call_lower(self, stmt: type[Statement], node: ast.Call) -> Result:
+        """Default lowering for Python call to statement.
+
+        This method is intended to be used by traits like `FromPythonCall` to
+        provide a default lowering for Python calls to statements.
+
+        Args:
+            stmt(type[Statement]): Statement class to construct.
+            node(ast.Call): Python call node to lower.
+
+        Returns:
+            Result: Result of lowering the Python call to statement.
+        """
+        args, kwargs = self.default_Call_inputs(stmt, node)
+        return Result(self.append_stmt(stmt(*args.values(), **kwargs)))
+
+    def default_Call_inputs(
+        self, stmt: type[Statement], node: ast.Call
+    ) -> tuple[dict[str, SSAValue | tuple[SSAValue, ...]], dict[str, Any]]:
+        from kirin.decl import fields
+        from kirin.dialects.py.data import PyAttr
+
+        fs = fields(stmt)
+        stmt_std_arg_names = fs.std_args.keys()
+        stmt_kw_args_name = fs.kw_args.keys()
+        stmt_attr_prop_names = fs.attr_or_props
+        stmt_required_names = fs.required_names
+        stmt_group_arg_names = fs.group_arg_names
+        args, kwargs = {}, {}
+        for name, value in zip(stmt_std_arg_names, node.args):
+            self._parse_arg(stmt_group_arg_names, args, name, value)
+        for kw in node.keywords:
+            if not isinstance(kw.arg, str):
+                raise DialectLoweringError("Expected string for keyword argument name")
+
+            arg: str = kw.arg
+            if arg in node.args:
+                raise DialectLoweringError(
+                    f"Keyword argument {arg} is already present in positional arguments"
+                )
+            elif arg in stmt_std_arg_names or arg in stmt_kw_args_name:
+                self._parse_arg(stmt_group_arg_names, kwargs, kw.arg, kw.value)
+            elif arg in stmt_attr_prop_names:
+                if not isinstance(kw.value, ast.Constant):
+                    raise DialectLoweringError(
+                        f"Expected constant for attribute or property {arg}"
+                    )
+                kwargs[arg] = PyAttr(kw.value.value)
+            else:
+                raise DialectLoweringError(f"Unexpected keyword argument {arg}")
+
+        for name in stmt_required_names:
+            if name not in args and name not in kwargs:
+                raise DialectLoweringError(f"Missing required argument {name}")
+
+        return args, kwargs
+
     def _parse_arg(
         self,
         group_names: set[str],

test/lowering/test_with.py

@@ -0,0 +1,38 @@
+from kirin import ir
+from kirin.decl import info, statement
+from kirin.prelude import python_no_opt
+from kirin.dialects import cf, py, func
+from kirin.lowering import Lowering
+
+dialect = ir.Dialect("test")
+
+
+@statement(dialect=dialect)
+class Adjoint(ir.Statement):
+    traits = frozenset({ir.FromPythonWithSingleItem()})
+    body: ir.Region = info.region()
+    result: ir.ResultValue = info.result()
+
+
+def with_example(x):
+    y = 1
+    with Adjoint() as f:  # type: ignore
+        y = x + 1
+    return y, f
+
+
+def test_with_lowering():
+    lowering = Lowering(python_no_opt.union([cf, func, dialect]))
+    code = lowering.run(with_example)
+    code.print()
+    assert isinstance(code, func.Function)
+    stmts = code.body.blocks[0].stmts
+    assert isinstance(stmts.at(0), py.Constant)
+    adjoint = stmts.at(1)
+    assert isinstance(adjoint, Adjoint)
+    assert len(adjoint.body.blocks) == 1
+    add = adjoint.body.blocks[0].stmts.at(1)
+    assert isinstance(add, py.Add)
+    assert isinstance(add.lhs, ir.BlockArgument)
+    assert isinstance(add.rhs, ir.SSAValue)
+    assert adjoint.result.name == "f"