Implement commute()

src/exo/API_scheduling.py

@@ -715,6 +715,39 @@ def reorder_stmts(proc, block_cursor):
     loopir  = Schedules.DoReorderStmt(loopir, s1, s2).result()
     return Procedure(loopir, _provenance_eq_Procedure=proc)
 
+@sched_op([ExprCursorA(many=True)])
+def commute_expr(proc, expr_cursors):
+    """
+    commute the binary operation of '+' and '*'.
+
+    args:
+        expr_cursors - a list of cursors to the binary operation
+
+    rewrite:
+        `a * b` <-- expr_cursor
+        -->
+        `b * a`
+
+        or
+
+        `a + b` <-- expr_cursor
+        -->
+        `b + a`
+    """
+
+    exprs   = [ ec._impl._node() for ec in expr_cursors ]
+    for e in exprs:
+        if not isinstance(e, LoopIR.BinOp) or (e.op != '+' and e.op != '*'):
+            raise TypeError(f"only '+' or '*' can commute, got {e.op}")
+    if any(not e.type.is_numeric() for e in exprs):
+        raise TypeError("only numeric (not index or size) expressions "
+                        "can commute by commute_expr()")
+
+    loopir  = proc._loopir_proc
+    loopir  = Schedules.DoCommuteExpr(loopir, exprs).result()
+    return Procedure(loopir, _provenance_eq_Procedure=proc)
+
+
 @sched_op([ExprCursorA(many=True), NameA, BoolA])
 def bind_expr(proc, expr_cursors, new_name, cse=False):
     """

src/exo/LoopIR_scheduling.py

@@ -1071,6 +1071,19 @@ def map_e(self, e):
             return super().map_e(e)
 
 
+class _DoCommuteExpr(LoopIR_Rewrite):
+    def __init__(self, proc, exprs):
+        self.exprs     = exprs
+        super().__init__(proc)
+        self.proc = InferEffects(self.proc).result()
+
+    def map_e(self, e):
+        if e in self.exprs:
+            assert isinstance(e, LoopIR.BinOp)
+            return e.update(lhs=e.rhs, rhs=e.lhs)
+        else:
+            return super().map_e(e)
+
 
 class _BindExpr(LoopIR_Rewrite):
     def __init__(self, proc, new_name, exprs, cse=False):
@@ -3643,3 +3656,4 @@ class Schedules:
     DoLiftAllocSimple  = _DoLiftAllocSimple
     DoFissionAfterSimple  = _DoFissionAfterSimple
     DoProductLoop  = _DoProductLoop
+    DoCommuteExpr  = _DoCommuteExpr

src/exo/stdlib/scheduling.py

@@ -22,6 +22,7 @@
     delete_pass,
     reorder_stmts,
     bind_expr,
+    commute_expr,
     #
     # subprocedure oriented operations
     extract_subproc,

tests/golden/test_schedules/test_commute.txt

@@ -0,0 +1,2 @@
+def foo(x: R[3] @ DRAM, y: R[3] @ DRAM, z: R @ DRAM):
+    z = y[2] * x[0]

tests/golden/test_schedules/test_commute3.txt

@@ -0,0 +1,2 @@
+def foo(x: R[3] @ DRAM, y: R[3] @ DRAM, z: R @ DRAM):
+    z = (x[1] + y[1] + y[2]) * (x[0] + y[0])

tests/test_schedules.py

@@ -7,6 +7,38 @@
 from exo import ParseFragmentError
 from exo.stdlib.scheduling import *
 
+def test_commute(golden):
+    @proc
+    def foo(x : R[3], y : R[3], z : R):
+        z = x[0]*y[2]
+    assert str(commute_expr(foo, 'x[0] * y[_]')) == golden
+
+def test_commute2():
+    @proc
+    def foo(x : R[3], y : R[3], z : R):
+        z = x[0] + y[0] + x[1] + y[1]
+
+    with pytest.raises(SchedulingError, match='failed to find matches'):
+        # TODO: Currently, expression pattern matching fails to find
+        # 'y[0]+x[1]' because LoopIR.BinOp is structured as (x[0], (y[0], (x[1], y[1]))).
+        # I think pattern matching should be powerful to find this.
+        commute_expr(foo, 'y[0] + x[1]')
+
+def test_commute3(golden):
+    @proc
+    def foo(x : R[3], y : R[3], z : R):
+        z = (x[0] + y[0]) * (x[1] + y[1] + y[2])
+    assert str(commute_expr(foo, '(x[_] + y[_]) * (x[_] + y[_] + y[_])')) == golden
+
+def test_commute4():
+    @proc
+    def foo(x : R[3], y : R[3], z : R):
+        z = x[0] - y[2]
+
+    with pytest.raises(TypeError, match="can commute"):
+        commute_expr(foo, 'x[0] - y[_]')
+
+
 def test_product_loop(golden):
     @proc
     def foo(n : size):