refactor scfs to use = yield

Author: makslevental

mlir_utils/ast/canonicalize.py

@@ -43,9 +43,8 @@ def transform_cst(
 
     logger.debug("[transformed code]\n\n%s", module_cst.code)
 
-    tree = ast.parse(module_cst.code, filename=inspect.getfile(f))
-    tree = ast.increment_lineno(tree, f.__code__.co_firstlineno - 1)
-    module_code_o = compile(tree, f.__code__.co_filename, "exec")
+    code = "\n" * (f.__code__.co_firstlineno - 1) + module_cst.code
+    module_code_o = compile(code, f.__code__.co_filename, "exec")
     new_f_code_o = next(
         c
         for c in module_code_o.co_consts

mlir_utils/dialects/ext/scf.py

@@ -77,16 +77,36 @@ def range_(
             yield iv, iter_args[0]
         else:
             yield iv
-    if len(iter_args):
-        previous_frame = inspect.currentframe().f_back
-        replacements = tuple(map(maybe_cast, for_op.results_))
-        _update_caller_vars(previous_frame, iter_args, replacements)
 
 
 def yield_(*args):
     if len(args) == 1 and isinstance(args[0], OpResultList):
         args = list(args[0])
-    yield__(args)
+    y = yield__(args)
+    parent_op = y.operation.parent.opview
+    if len(parent_op.results_):
+        results = get_result_or_results(parent_op)
+        assert (
+            isinstance(results, (OpResult, OpResultList))
+            or isinstance(results, list)
+            and all(isinstance(r, OpResult) for r in results)
+        ), f"api has changed: {results=}"
+        if isinstance(results, OpResult):
+            results = [results]
+        unpacked_args = args
+        if any(isinstance(a, OpResultList) for a in unpacked_args):
+            assert len(unpacked_args) == 1
+            unpacked_args = list(unpacked_args[0])
+
+        assert len(results) == len(unpacked_args), f"{results=}, {unpacked_args=}"
+        for i, r in enumerate(results):
+            if r.type == T._placeholder_opaque_t():
+                r.set_type(unpacked_args[i].type)
+
+        results = maybe_cast(results)
+        if len(results) > 1:
+            return results
+        return results[0]
 
 
 def _if(cond, results_=None, *, has_else=False, loc=None, ip=None):
@@ -184,22 +204,51 @@ def stack_if(*args, **kwargs):
     return IfStack.push(*args, **kwargs)
 
 
-def stack_yield(*args):
-    return IfStack.yield_(*args)
+def unstack_if(cond: Value, results_=None, has_else=False):
+    if results_ is None:
+        results_ = []
+    if results_:
+        has_else = True
+    assert isinstance(cond, Value), f"cond must be a mlir.Value: {cond=}"
+    if_op = _if(cond, results_, has_else=has_else)
+    cond.owner.move_before(if_op)
+
+    ip = InsertionPoint(if_op.then_block)
+    ip.__enter__()
+
+    return ip, if_op
 
 
 def end_branch():
     IfStack.pop_branch()
 
 
+def unstack_end_branch(ip):
+    ip.__exit__(None, None, None)
+
+
 def else_():
     return IfStack.push_else()
 
 
+def unstack_else_if(if_op):
+    assert len(
+        if_op.regions[1].blocks
+    ), f"can't have else without bb in second region of {if_op=}"
+
+    ip = InsertionPoint(if_op.else_block)
+    ip.__enter__()
+    return ip
+
+
 def end_if():
     IfStack.pop()
 
 
+def unstack_end_if(ip):
+    ip.__exit__(None, None, None)
+
+
 def insert_body_maybe_semicolon(
     node: cst.CSTNode, index: int, new_node: cst.CSTNode, before=False
 ):
@@ -233,125 +282,121 @@ def insert_body_maybe_semicolon(
 
 
 class ReplaceYieldWithSCFYield(StrictTransformer):
-    @m.call_if_inside(m.If(test=m.NamedExpr(value=m.Comparison())))
+    @m.call_if_inside(m.If())
     @m.leave(m.Yield(value=m.Tuple()))
     def tuple_yield_inside_conditional(
-        self, original_node: cst.Yield, updated_node: cst.Yield
+        self, original_node: cst.Yield, _updated_node: cst.Yield
     ):
         args = [cst.Arg(e.value) for e in original_node.value.elements]
-        return ast_call(stack_yield.__name__, args)
+        return ast_call(yield_.__name__, args)
 
-    @m.call_if_inside(m.If(test=m.NamedExpr(value=m.Comparison())))
+    @m.call_if_inside(m.If())
     @m.leave(m.Yield(value=~m.Tuple()))
     def single_yield_inside_conditional(
-        self, original_node: cst.Yield, updated_node: cst.Yield
+        self, original_node: cst.Yield, _updated_node: cst.Yield
     ):
         args = [cst.Arg(original_node.value)] if original_node.value else []
-        return ast_call(stack_yield.__name__, args)
+        return ast_call(yield_.__name__, args)
 
-    @m.call_if_not_inside(m.If(test=m.NamedExpr(value=m.Comparison())))
+    @m.call_if_not_inside(m.If())
     @m.leave(m.Yield(value=m.Tuple()))
-    def tuple_yield(self, original_node: cst.Yield, updated_node: cst.Yield):
+    def tuple_yield(self, original_node: cst.Yield, _updated_node: cst.Yield):
         args = [cst.Arg(e.value) for e in original_node.value.elements]
         return ast_call(yield_.__name__, args)
 
-    @m.call_if_not_inside(m.If(test=m.NamedExpr(value=m.Comparison())))
+    @m.call_if_not_inside(m.If())
     @m.leave(m.Yield(value=~m.Tuple()))
-    def single_yield(self, original_node: cst.Yield, updated_node: cst.Yield):
+    def single_yield(self, original_node: cst.Yield, _updated_node: cst.Yield):
         args = [cst.Arg(original_node.value)] if original_node.value else []
         return ast_call(yield_.__name__, args)
 
 
-class InsertEmptySCFYield(StrictTransformer):
+class InsertEmptyYield(StrictTransformer):
     @m.leave(m.If() | m.Else())
     def leave_(
         self, _original_node: cst.If | cst.Else, updated_node: cst.If | cst.Else
     ) -> cst.If | cst.Else:
         indented_block = updated_node.body
         last_statement = indented_block.body[-1]
-        if not m.matches(last_statement, m.SimpleStatementLine([m.Expr(m.Yield())])):
-            return insert_body_maybe_semicolon(
-                updated_node, -1, ast_call(yield_.__name__)
-            )
+        if not m.matches(last_statement, m.SimpleStatementLine()):
+            return insert_body_maybe_semicolon(updated_node, -1, cst.Yield())
+        elif m.matches(last_statement, m.SimpleStatementLine()) and not m.findall(
+            last_statement, m.Yield()
+        ):
+            return insert_body_maybe_semicolon(updated_node, -1, cst.Yield())
         # VERY IMPORTANT: you have to return the updated node if you believe
         # at any point there was a mutation anywhere in the tree below
         return updated_node
 
 
 class CanonicalizeElIfs(StrictTransformer):
-    @m.leave(m.If(orelse=m.If(test=m.NamedExpr())))
-    def leave_if_with_elif_named(
+    @m.leave(m.If(orelse=m.If()))
+    def leave_if_with_elif(
         self, _original_node: cst.If, updated_node: cst.If
     ) -> cst.If:
-        return updated_node.with_changes(
-            orelse=cst.Else(
-                cst.IndentedBlock(
-                    [
-                        updated_node.orelse,
-                        cst.SimpleStatementLine(
-                            [cst.Expr(cst.Yield(updated_node.orelse.test.target))]
+        indented_block = updated_node.orelse.body
+        last_statement = indented_block.body[-1]
+        if m.matches(last_statement, m.SimpleStatementLine()) and m.matches(
+            last_statement.body[-1], m.Assign(value=m.Yield())
+        ):
+            assign_targets = last_statement.body[-1].targets
+            last_statement = cst.SimpleStatementLine(
+                [
+                    cst.Assign(
+                        targets=assign_targets,
+                        value=cst.Yield(
+                            cst.Tuple([cst.Element(a.target) for a in assign_targets])
+                            if len(assign_targets) > 1
+                            else assign_targets[0].target
                         ),
-                    ]
-                )
+                    )
+                ]
             )
-        )
-
-    @m.leave(m.If(orelse=m.If(test=~m.NamedExpr())))
-    def leave_if_with_elif(
-        self, _original_node: cst.If, updated_node: cst.If
-    ) -> cst.If:
-        return updated_node.with_changes(
-            orelse=cst.Else(cst.IndentedBlock([updated_node.orelse]))
-        )
+            body = [updated_node.orelse, last_statement]
+        else:
+            body = [updated_node.orelse]
+        return updated_node.with_changes(orelse=cst.Else(cst.IndentedBlock(body)))
 
 
 class ReplaceSCFCond(StrictTransformer):
-    @m.leave(m.If(test=m.NamedExpr(value=m.Call(func=m.Name(stack_if.__name__)))))
+    @m.leave(m.If(test=m.Call(func=m.Name(stack_if.__name__))))
     def insert_with_results(
         self, original_node: cst.If, _updated_node: cst.If
     ) -> cst.If:
         return original_node
 
-    @m.leave(m.If(test=m.NamedExpr(value=m.Comparison())))
+    @m.leave(m.If(test=~m.Call(func=m.Name(stack_if.__name__))))
     def insert_with_results(
         self, original_node: cst.If, updated_node: cst.If
     ) -> cst.If:
         indented_block = updated_node.body
         last_statement = indented_block.body[-1]
         assert m.matches(
             last_statement, m.SimpleStatementLine()
-        ), f"conditional with := must explicitly yield on last line"
-        yield_expr = last_statement.body[0]
-        if m.matches(yield_expr.value, m.Call(func=m.Name(stack_yield.__name__))):
-            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))] * len(
-                yield_expr.value.args
-            )
-        elif m.matches(yield_expr.value.value, m.Name()):
-            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))]
-        elif m.matches(yield_expr.value.value, m.Tuple()):
-            results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))] * len(
-                yield_expr.value.value.elements
-            )
+        ), f"conditional must end with a statement"
+        yield_expr = m.findall(last_statement, m.Call(func=m.Name(yield_.__name__)))
+        assert (
+            len(yield_expr) == 1
+        ), f"conditional with must explicitly {yield_.__name__} on last line: {yield_expr}"
+        yield_expr = yield_expr[0]
+        results = [cst.Element(ast_call(T._placeholder_opaque_t.__name__))] * len(
+            yield_expr.args
+        )
         results = cst.Tuple(results)
 
         test = original_node.test
-        compare = test.value
-        assert m.matches(
-            compare, m.Comparison()
-        ), f"expected cst.Compare from {compare=}"
-        new_compare = ast_call(
-            stack_if.__name__, args=[cst.Arg(compare), cst.Arg(results)]
+        new_test = ast_call(
+            stack_if.__name__,
+            args=[
+                cst.Arg(test),
+                cst.Arg(results),
+                cst.Arg(
+                    cst.Name(str(bool(original_node.orelse))),
+                    keyword=cst.Name("has_else"),
+                ),
+            ],
         )
-        new_test = test.deep_replace(compare, new_compare)
-        return updated_node.with_changes(test=new_test)
-
-    @m.leave(m.If(test=m.Comparison()))
-    def insert_no_results(self, original_node: cst.If, updated_node: cst.If) -> cst.If:
-        test = original_node.test
-        args = [cst.Arg(test)]
-        if original_node.orelse:
-            args += [cst.Arg(cst.Tuple([])), cst.Arg(cst.Name(str(True)))]
-        new_test = ast_call(stack_if.__name__, args=args)
+        new_test = test.deep_replace(test, new_test)
         return updated_node.with_changes(test=new_test)
 
 
@@ -424,7 +469,6 @@ def patch_bytecode(self, code: ConcreteBytecode, f):
         f.__globals__[end_branch.__name__] = end_branch
         f.__globals__[end_if.__name__] = end_if
         f.__globals__[stack_if.__name__] = stack_if
-        f.__globals__[stack_yield.__name__] = stack_yield
         f.__globals__[yield_.__name__] = yield_
         f.__globals__[T._placeholder_opaque_t.__name__] = T._placeholder_opaque_t
         return code
@@ -433,7 +477,7 @@ def patch_bytecode(self, code: ConcreteBytecode, f):
 class SCFCanonicalizer(Canonicalizer):
     cst_transformers = [
         CanonicalizeElIfs,
-        InsertEmptySCFYield,
+        InsertEmptyYield,
         ReplaceYieldWithSCFYield,
         ReplaceSCFCond,
         InsertEndIfs,

mlir_utils/util.py

@@ -36,7 +36,7 @@ def get_result_or_results(
     if isinstance(op, Value):
         return op
     return (
-        get_op_results_or_values(op)
+        list(get_op_results_or_values(op))
         if len(op.operation.results) > 1
         else get_op_result_or_value(op)
         if len(op.operation.results) > 0
@@ -94,14 +94,17 @@ def get_value_caster(typeid: TypeID):
     return __VALUE_CASTERS[typeid]
 
 
-def maybe_cast(val: Value):
+def maybe_cast(val: Value | list[Value]):
     """Maybe cast an ir.Value to one of Tensor, Scalar.
 
     Args:
       val: The ir.Value to maybe cast.
     """
     from mlir_utils.dialects.ext.arith import Scalar
 
+    if isinstance(val, list):
+        return list(map(maybe_cast, val))
+
     if not isinstance(val, Value):
         return val
 

tests/test_location_tracking.py

@@ -34,12 +34,12 @@ def test_if_replace_yield_5(ctx: MLIRContext):
     def iffoo():
         one = constant(1.0)
         two = constant(2.0)
-        if res := stack_if(one < two, (T.f64_t, T.f64_t, T.f64_t)):
+        if stack_if(one < two, (T.f64_t, T.f64_t, T.f64_t)):
             three = constant(3.0)
-            yield three, three, three
+            res1, res2, res3 = yield three, three, three
         else:
             four = constant(4.0)
-            yield four, four, four
+            res1, res2, res3 = yield four, four, four
         return
 
     iffoo()
@@ -49,24 +49,24 @@ def iffoo():
     module {{
       %cst = arith.constant 1.000000e+00 : f64 THIS_DIR{sep}test_location_tracking.py:35:10
       %cst_0 = arith.constant 2.000000e+00 : f64 THIS_DIR{sep}test_location_tracking.py:36:10
-      %0 = arith.cmpf olt, %cst, %cst_0 : f64 THIS_DIR{sep}test_location_tracking.py:37:23
+      %0 = arith.cmpf olt, %cst, %cst_0 : f64 THIS_DIR{sep}test_location_tracking.py:37:16
       %1:3 = scf.if %0 -> (f64, f64, f64) {{
         %cst_1 = arith.constant 3.000000e+00 : f64 THIS_DIR{sep}test_location_tracking.py:38:16
-        scf.yield %cst_1, %cst_1, %cst_1 : f64, f64, f64 THIS_DIR{sep}test_location_tracking.py:39:8
+        scf.yield %cst_1, %cst_1, %cst_1 : f64, f64, f64 THIS_DIR{sep}test_location_tracking.py:39:27
       }} else {{
         %cst_1 = arith.constant 4.000000e+00 : f64 THIS_DIR{sep}test_location_tracking.py:41:24
-        scf.yield %cst_1, %cst_1, %cst_1 : f64, f64, f64 THIS_DIR{sep}test_location_tracking.py:42:8
-      }} THIS_DIR{sep}test_location_tracking.py:37:14
+        scf.yield %cst_1, %cst_1, %cst_1 : f64, f64, f64 THIS_DIR{sep}test_location_tracking.py:42:27
+      }} THIS_DIR{sep}test_location_tracking.py:37:7
     }} [unknown]
     #loc = [unknown]
     #loc1 = THIS_DIR{sep}test_location_tracking.py:35:10
     #loc2 = THIS_DIR{sep}test_location_tracking.py:36:10
-    #loc3 = THIS_DIR{sep}test_location_tracking.py:37:23
-    #loc4 = THIS_DIR{sep}test_location_tracking.py:37:14
+    #loc3 = THIS_DIR{sep}test_location_tracking.py:37:16
+    #loc4 = THIS_DIR{sep}test_location_tracking.py:37:7
     #loc5 = THIS_DIR{sep}test_location_tracking.py:38:16
-    #loc6 = THIS_DIR{sep}test_location_tracking.py:39:8
+    #loc6 = THIS_DIR{sep}test_location_tracking.py:39:27
     #loc7 = THIS_DIR{sep}test_location_tracking.py:41:24
-    #loc8 = THIS_DIR{sep}test_location_tracking.py:42:8
+    #loc8 = THIS_DIR{sep}test_location_tracking.py:42:27
     """
     )
     asm = get_asm(ctx.module.operation)