transformations: (licm) add can_be_hoisted helper function

tests/transforms/test_loop_invariant_code_motion.py

@@ -0,0 +1,77 @@
+from xdsl.builder import ImplicitBuilder
+from xdsl.dialects import test
+from xdsl.dialects.builtin import IndexType
+from xdsl.ir import Block, Region
+from xdsl.transforms.loop_invariant_code_motion import can_be_hoisted
+
+index = IndexType()
+
+
+def test_can_be_hoisted():
+    """Test for with loop-variant variables"""
+
+    outer_region = Region(Block(arg_types=(index,)))
+    mid_region = Region(Block(arg_types=(index,)))
+    inner_region = Region(Block(arg_types=(index,)))
+
+    with ImplicitBuilder(outer_region) as (outer_arg,):
+        outer_op_no_args = test.TestOp()
+        outer_op_outer_arg = test.TestOp(operands=(outer_arg,))
+        with ImplicitBuilder(mid_region) as (mid_arg,):
+            mid_op_no_args = test.TestOp()
+            mid_op_outer_arg = test.TestOp(operands=(outer_arg,))
+            mid_op_mid_arg = test.TestOp(operands=(mid_arg,))
+            with ImplicitBuilder(inner_region) as (inner_arg,):
+                inner_op_no_args = test.TestOp()
+                inner_op_outer_arg = test.TestOp(operands=(outer_arg,))
+                inner_op_inner_arg = test.TestOp(operands=(inner_arg,))
+                inner_term_op = test.TestTermOp()
+            mid_term_op = test.TestTermOp()
+            inner_region_op = test.TestOp(regions=(inner_region,))
+        mid_region_op = test.TestOp(regions=(mid_region,))
+        out_term_op = test.TestTermOp()
+
+    # outer region
+    assert can_be_hoisted(outer_op_no_args, outer_region)
+    assert not can_be_hoisted(outer_op_outer_arg, outer_region)
+    assert can_be_hoisted(mid_op_no_args, outer_region)
+    assert not can_be_hoisted(mid_op_outer_arg, outer_region)
+    assert not can_be_hoisted(mid_op_mid_arg, outer_region)
+    assert can_be_hoisted(inner_op_no_args, outer_region)
+    assert not can_be_hoisted(inner_op_outer_arg, outer_region)
+    assert not can_be_hoisted(inner_op_inner_arg, outer_region)
+    assert not can_be_hoisted(inner_term_op, outer_region)
+    assert not can_be_hoisted(mid_term_op, outer_region)
+    assert not can_be_hoisted(inner_region_op, outer_region)
+    assert not can_be_hoisted(mid_region_op, outer_region)
+    assert not can_be_hoisted(out_term_op, outer_region)
+
+    # mid region
+    assert can_be_hoisted(outer_op_no_args, mid_region)
+    assert can_be_hoisted(outer_op_outer_arg, mid_region)
+    assert can_be_hoisted(mid_op_no_args, mid_region)
+    assert can_be_hoisted(mid_op_outer_arg, mid_region)
+    assert not can_be_hoisted(mid_op_mid_arg, mid_region)
+    assert can_be_hoisted(inner_op_no_args, mid_region)
+    assert can_be_hoisted(inner_op_outer_arg, mid_region)
+    assert not can_be_hoisted(inner_op_inner_arg, mid_region)
+    assert not can_be_hoisted(inner_term_op, mid_region)
+    assert not can_be_hoisted(mid_term_op, mid_region)
+    assert can_be_hoisted(inner_region_op, mid_region)
+    assert can_be_hoisted(mid_region_op, mid_region)
+    assert not can_be_hoisted(out_term_op, mid_region)
+
+    # inner region
+    assert can_be_hoisted(outer_op_no_args, inner_region)
+    assert can_be_hoisted(outer_op_outer_arg, inner_region)
+    assert can_be_hoisted(mid_op_no_args, inner_region)
+    assert can_be_hoisted(mid_op_outer_arg, inner_region)
+    assert can_be_hoisted(mid_op_mid_arg, inner_region)
+    assert can_be_hoisted(inner_op_no_args, inner_region)
+    assert can_be_hoisted(inner_op_outer_arg, inner_region)
+    assert not can_be_hoisted(inner_op_inner_arg, inner_region)
+    assert not can_be_hoisted(inner_term_op, inner_region)
+    assert not can_be_hoisted(mid_term_op, inner_region)
+    assert can_be_hoisted(inner_region_op, inner_region)
+    assert can_be_hoisted(mid_region_op, inner_region)
+    assert not can_be_hoisted(out_term_op, inner_region)

xdsl/transforms/loop_invariant_code_motion.py

@@ -0,0 +1,42 @@
+"""
+This pass hoists operation that are invariant to the loops.
+
+Similar to MLIR's loop invariant code motion:
+
+https://mlir.llvm.org/doxygen/LoopInvariantCodeMotion_8cpp_source.html
+
+An operation is loop-invariant if it depends only of values defined outside of the loop.
+LICM moves these operations out of the loop body so that they are not computed more than
+once.
+
+  for i in range(x, N, M):                for i in range(x, N, M):
+    for j in range(0, M, K):    ---->        c[i]= A[1] + b[1]
+      c[i]=A[1]+b[1]
+"""
+
+from xdsl.ir import Operation, Region
+from xdsl.traits import IsTerminator
+
+
+def can_be_hoisted(op: Operation, target_region: Region) -> bool | None:
+    """
+    Checks whether the given op can be hoisted by checking that
+    - the op and none of its contained operations depend on values inside of the
+     loop.
+    """
+    #  Do not move terminators.
+    if op.has_trait(IsTerminator):
+        return False
+
+    # Walk the nested operations and check that all used values are either
+    # defined outside of the loop or in a nested region, but not at the level of
+    # the loop body.
+    for child in op.walk():
+        for operand in child.operands:
+            operand_owner = operand.owner
+            assert operand_owner is not None
+            if op.is_ancestor(operand_owner):
+                continue
+            if target_region.is_ancestor(operand_owner):
+                return False
+    return True