Repository: exo-lang/exo

Issue:
Implement merge_reduce
Open to suggestions for a better name
```
loc = expr1
loc += expr2
--->
loc = expr1 + expr2
```

Discussion:
yamaguchi1024: Looks good to me. @gilbo @alexreinking Do you want to comment?
alexreinking: This is a useful change, but I think it would be a _lot_ more useful if it were generalized slightly. I think this should be something more like `fuse_writes` which takes two adjacent writes to the same location and condenses them down to a single write. There are four cases:

1. Write-Write. The second write wins, the first write is deleted.
2. Write-Accum. This PR.
3. Accum-Write. The write wins, the accum is deleted.
4. Accum-Accum. `x += y ; x += z` -> `x += y + z`.

Then we can implement a generalization to `n` consecutive writes on top of this.
yamaguchi1024: @alexreinking Generalizing this more sounds useful/reasonable, @skeqiqevian can you take a look?
alexreinking: @skeqiqevian please push an empty commit to make CI run.
```
git commit --allow-empty -m "trigger CI"
```
To avoid this in the future, please run:
```
python -m pip install -r requirements.txt
pre-commit install
```
Then the code formatter will run whenever you commit.


Patch:
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index 0abc514..dcdb473 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -1592,6 +1592,63 @@ def reorder_loops(proc, nested_loops):
     return Procedure(loopir, _provenance_eq_Procedure=proc)
 
 
+@sched_op([BlockCursorA(block_size=2)])
+def merge_writes(proc, block_cursor):
+    """
+    Merge consecutive assign and reduce statement into a single statement.
+    Handles all 4 cases of (assign, reduce) x (reduce, assign).
+
+    args:
+        block_cursor          - cursor pointing to the block of two consecutive
+                                assign/reduce statement.
+
+    rewrite:
+        `a = b`
+        `a = c`
+            ->
+        `a = c`
+        ----------------------
+        `a += b`
+        `a = c`
+            ->
+        `a = c`
+        ----------------------
+        `a = b`
+        `a += c`
+            ->
+        `a = b + c`
+        ----------------------
+        `a += b`
+        `a += c`
+            ->
+        `a += b + c`
+        ----------------------
+
+    """
+    stmt1 = block_cursor[0]._impl._node()
+    stmt2 = block_cursor[1]._impl._node()
+
+    if not isinstance(stmt1, (LoopIR.Assign, LoopIR.Reduce)) or not isinstance(
+        stmt2, (LoopIR.Assign, LoopIR.Reduce)
+    ):
+        raise ValueError(
+            f"expected two consecutive assign/reduce statements, "
+            f"got {type(stmt1)} and {type(stmt2)} instead."
+        )
+    if stmt1.name != stmt2.name or stmt1.type != stmt2.type:
+        raise ValueError(
+            "expected the two statements' left hand sides to have the same name & type"
+        )
+    if stmt1.rhs.type != stmt2.rhs.type:
+        raise ValueError(
+            "expected the two statements' right hand sides to have the same type."
+        )
+
+    loopir = proc._loopir_proc
+    loopir = Schedules.DoMergeWrites(loopir, stmt1, stmt2).result()
+    return Procedure(loopir, _provenance_eq_Procedure=proc)
+
+
 @sched_op([GapCursorA, PosIntA])
 def fission(proc, gap_cursor, n_lifts=1):
     """
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index a163101..b5cdeec 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -226,6 +226,64 @@ class _DoProductLoop(LoopIR_Rewrite):
         return super().map_e(e)
 
 
+def get_reads(e):
+    if isinstance(e, LoopIR.Read):
+        return sum([get_reads(e) for e in e.idx], [(e.name, e.type)])
+    elif isinstance(e, LoopIR.USub):
+        return get_reads(e.arg)
+    elif isinstance(e, LoopIR.BinOp):
+        return get_reads(e.lhs) + get_reads(e.rhs)
+    elif isinstance(e, LoopIR.BuiltIn):
+        return sum([get_reads(a) for a in e.args], [])
+    elif isinstance(e, LoopIR.Const):
+        return []
+    else:
+        assert False, "bad case"
+
+
+class _DoMergeWrites(LoopIR_Rewrite):
+    def __init__(self, proc, stmt1, stmt2):
+        try:
+            assert len(stmt1.idx) == len(stmt2.idx)
+            for i, j in zip(stmt1.idx, stmt2.idx):
+                Check_ExprEqvInContext(proc, [stmt1, stmt2], i, j)
+        except SchedulingError as e:
+            raise SchedulingError(
+                "expected the left hand side's indices to be the same."
+            ) from e
+
+        if any(
+            stmt1.name == name and stmt1.type == typ
+            for name, typ in get_reads(stmt2.rhs)
+        ):
+            raise SchedulingError(
+                "expected the right hand side of the second statement to not "
+                "depend on the left hand side of the first statement."
+            )
+
+        self.new_rhs = LoopIR.BinOp(
+            "+", stmt1.rhs, stmt2.rhs, stmt1.type, stmt1.srcinfo
+        )
+        self.s1 = stmt1
+        self.s2 = stmt2
+
+        super().__init__(proc)
+
+        self.proc = InferEffects(self.proc).result()
+
+    def map_stmts(self, stmts):
+        if isinstance(self.s2, LoopIR.Assign):
+            for i, s in enumerate(stmts):
+                if s is self.s2:
+                    return stmts[: i - 1] + stmts[i:]
+        else:
+            for i, s in enumerate(stmts):
+                if s is self.s1:
+                    return stmts[:i] + [s.update(rhs=self.new_rhs)] + stmts[i + 2 :]
+
+        return super().map_stmts(stmts)
+
+
 class _Reorder(LoopIR_Rewrite):
     def __init__(self, proc, loop_stmt):
         self.stmt = loop_stmt
@@ -3909,3 +3967,4 @@ class Schedules:
     DoFissionAfterSimple = _DoFissionAfterSimple
     DoProductLoop = _DoProductLoop
     DoCommuteExpr = _DoCommuteExpr
+    DoMergeWrites = _DoMergeWrites
diff --git a/src/exo/stdlib/scheduling.py b/src/exo/stdlib/scheduling.py
index ccc8040..1043a8b 100644
--- a/src/exo/stdlib/scheduling.py
+++ b/src/exo/stdlib/scheduling.py
@@ -56,6 +56,7 @@ from ..API_scheduling import (
     mult_loops,
     cut_loop,
     reorder_loops,
+    merge_writes,
     fission,
     fusion,
     remove_loop,


Test patch:
diff --git a/tests/golden/test_schedules/test_merge_writes_all_4_cases.txt b/tests/golden/test_schedules/test_merge_writes_all_4_cases.txt
new file mode 100644
index 0000000..f9ac7e7
--- /dev/null
+++ b/tests/golden/test_schedules/test_merge_writes_all_4_cases.txt
@@ -0,0 +1,8 @@
+def bar(x: R[4] @ DRAM, y: R[4] @ DRAM):
+    for i in seq(0, 10):
+        if i < 5:
+            tmp: R[4] @ DRAM
+            tmp[0] = y[0]
+            tmp[1] = x[1] + y[1]
+            tmp[2] = y[2]
+            tmp[3] += x[3] + y[3]
diff --git a/tests/golden/test_schedules/test_merge_writes_array_indexing.txt b/tests/golden/test_schedules/test_merge_writes_array_indexing.txt
new file mode 100644
index 0000000..f8632a5
--- /dev/null
+++ b/tests/golden/test_schedules/test_merge_writes_array_indexing.txt
@@ -0,0 +1,6 @@
+def bar(x: R[3] @ DRAM, y: R[3] @ DRAM, z: R @ DRAM):
+    for i in seq(0, 3):
+        for j in seq(0, 3):
+            if i < 2:
+                tmp: R[4, 4] @ DRAM
+                tmp[i + j, j] = x[i] + y[j]
diff --git a/tests/golden/test_schedules/test_merge_writes_consecutively.txt b/tests/golden/test_schedules/test_merge_writes_consecutively.txt
new file mode 100644
index 0000000..2f349c3
--- /dev/null
+++ b/tests/golden/test_schedules/test_merge_writes_consecutively.txt
@@ -0,0 +1,3 @@
+def bar(w: R @ DRAM, x: R @ DRAM, y: R @ DRAM, z: R @ DRAM):
+    z = w + x + y
+    w = x
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 9e9cde7..6cac342 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -624,6 +624,122 @@ def test_simple_reorder(golden):
     assert str(bar) == golden
 
 
+def test_merge_writes_all_4_cases(golden):
+    @proc
+    def bar(x: R[4], y: R[4]):
+        for i in seq(0, 10):
+            if i < 5:
+                tmp: R[4]
+                tmp[0] = x[0]
+                tmp[0] = y[0]
+                tmp[1] = x[1]
+                tmp[1] += y[1]
+                tmp[2] += x[2]
+                tmp[2] = y[2]
+                tmp[3] += x[3]
+                tmp[3] += y[3]
+
+    bar = merge_writes(bar, "tmp[0] = x[0]; tmp[0] = y[0]")
+    bar = merge_writes(bar, "tmp[1] = x[1]; tmp[1] += y[1]")
+    bar = merge_writes(bar, "tmp[2] += x[2]; tmp[2] = y[2]")
+    bar = merge_writes(bar, "tmp[3] += x[3]; tmp[3] += y[3]")
+    assert str(bar) == golden
+
+
+def test_merge_writes_consecutively(golden):
+    @proc
+    def bar(w: R, x: R, y: R, z: R):
+        z = w
+        z += x
+        z += y
+        w = x
+
+    bar = merge_writes(bar, "z = w; z += x")
+    bar = merge_writes(bar, "z = w + x; z += y")
+    assert str(bar) == golden
+
+
+def test_merge_writes_array_indexing(golden):
+    @proc
+    def bar(x: R[3], y: R[3], z: R):
+        for i in seq(0, 3):
+            for j in seq(0, 3):
+                if i < 2:
+                    tmp: R[4, 4]
+                    tmp[i + j, j] = x[i]
+                    tmp[i + j, j] += y[j]
+
+    bar = merge_writes(bar, "tmp[i+j, j] = x[i]; tmp[i+j, j] += y[j]")
+    assert str(bar) == golden
+
+
+def test_merge_writes_second_rhs_depends_on_first_lhs(golden):
+    @proc
+    def bar(x: R[5], y: R[3]):
+        for i in seq(0, 3):
+            if i > 0:
+                x[2 * i - 1] = x[i] + y[i]
+                x[2 * i - 1] += x[i]
+
+    with pytest.raises(
+        SchedulingError, match="expected the right hand side of the second statement"
+    ):
+        bar = merge_writes(bar, "x[2*i-1] = x[i] + y[i]; x[2*i-1] += x[i]")
+
+
+def test_merge_writes_wrong_type_error(golden):
+    @proc
+    def bar(x: R, y: R):
+        for i in seq(0, 10):
+            y = x
+            if i > 5:
+                y = x
+
+    with pytest.raises(
+        ValueError, match="expected two consecutive assign/reduce statements"
+    ):
+        bar = merge_writes(bar, "y = x; _")
+
+
+def test_merge_writes_different_lhs_error(golden):
+    @proc
+    def bar(x: R, y: R):
+        x = y
+        y += x
+
+    with pytest.raises(
+        ValueError,
+        match="expected the two statements' left hand sides to have the same name & type",
+    ):
+        bar = merge_writes(bar, "x = y; y += x")
+
+
+def test_merge_writes_different_lhs_arrays_error(golden):
+    @proc
+    def bar(x: R[3], y: R):
+        x[0] = y
+        x[1] += y
+
+    with pytest.raises(
+        SchedulingError, match="expected the left hand side's indices to be the same."
+    ):
+        bar = merge_writes(bar, "x[0] = y; x[1] += y")
+
+
+def test_merge_writes_different_lhs_arrays_error(golden):
+    @proc
+    def bar(x: R[3, 3], y: R):
+        z = x[0:2, 0:2]
+        for i in seq(0, 2):
+            z[i, 1] = y
+            z[i + 1, 1] += y
+
+    with pytest.raises(
+        SchedulingError, match="expected the left hand side's indices to be the same."
+    ):
+        bar = merge_writes(bar, "z[i, 1] = y; z[i+1, 1] += y")
+
+
 def test_simple_unroll(golden):
     @proc
     def bar(A: i8[10]):