Repository: exo-lang/exo

Issue:
Support unroll_buffer

Discussion:
yamaguchi1024: The CI is failing due to pip failure. Will trigger the build again tomorrow
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
> Merging [#460](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (854e44b) into [master](https://app.codecov.io/gh/exo-lang/exo/commit/acc319250855c8af88c06499754bd599769bc01d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (acc3192) will **increase** coverage by `0.02%`.
> The diff coverage is `91.83%`.

> :exclamation: Current head 854e44b differs from pull request most recent head 4e88474. Consider uploading reports for the commit 4e88474 to get more accurate results

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang).

```diff
@@            Coverage Diff             @@
##           master     #460      +/-   ##
==========================================
+ Coverage   86.49%   86.52%   +0.02%     
==========================================
  Files          78       78              
  Lines       17796    17894      +98     
==========================================
+ Hits        15392    15482      +90     
- Misses       2404     2412       +8     
```


| [Impacted Files](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) | Coverage Δ | |
|---|---|---|
| [src/exo/stdlib/scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9zdGRsaWIvc2NoZWR1bGluZy5weQ==) | `83.65% <ø> (ø)` | |
| [src/exo/LoopIR\_scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfc2NoZWR1bGluZy5weQ==) | `90.64% <89.87%> (-0.03%)` | :arrow_down: |
| [src/exo/API\_scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9BUElfc2NoZWR1bGluZy5weQ==) | `88.81% <100.00%> (+0.06%)` | :arrow_up: |
| [tests/test\_schedules.py](https://app.codecov.io/gh/exo-lang/exo/pull/460?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-dGVzdHMvdGVzdF9zY2hlZHVsZXMucHk=) | `99.36% <100.00%> (+0.01%)` | :arrow_up: |



Patch:
diff --git a/src/exo/API_scheduling.py b/src/exo/API_scheduling.py
index 3ffd66a..30fb98e 100644
--- a/src/exo/API_scheduling.py
+++ b/src/exo/API_scheduling.py
@@ -1268,6 +1268,28 @@ def mult_dim(proc, alloc_cursor, hi_dim_idx, lo_dim_idx):
     return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
 
 
+@sched_op([AllocCursorA, IntA])
+def unroll_buffer(proc, alloc_cursor, dimension):
+    """
+    Unroll the buffer allocation with constant dimension.
+
+    args:
+        alloc_cursor  - cursor to the buffer with constant dimension
+        dimension     - dimension to unroll
+
+    rewrite:
+        `buf : T[2]` <- alloc_cursor
+        `...`
+        ->
+        `buf_0 : T`
+        `buf_1 : T`
+        `...`
+    """
+
+    ir, fwd = scheduling.DoUnrollBuffer(alloc_cursor._impl, dimension)
+    return Procedure(ir, _provenance_eq_Procedure=proc, _forward=fwd)
+
+
 @sched_op([AllocCursorA, PosIntA])
 def lift_alloc(proc, alloc_cursor, n_lifts=1):
     """
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index 3459d79..1b5d585 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -3614,6 +3614,104 @@ class DoBoundAlloc(Cursor_Rewrite):
         return new_stmts
 
 
+def DoUnrollBuffer(alloc_cursor, dim):
+    alloc_stmt = alloc_cursor._node
+
+    assert isinstance(alloc_stmt, LoopIR.Alloc)
+
+    if not alloc_stmt.type.shape():
+        raise SchedulingError("Cannot unroll a scalar buffer")
+
+    buf_size = alloc_stmt.type.shape()[dim]
+    if not isinstance(buf_size, LoopIR.Const):
+        raise SchedulingError(
+            f"Expected a constant buffer dimension, got {buf_size} at {dim}'th dimension."
+        )
+
+    used_allocs = set()
+    buf_syms = []
+    for i in range(0, buf_size.val):
+        new_name = str(alloc_stmt.name) + "_" + str(i)
+        buf_syms.append(Sym(new_name))
+
+    def mk_read(c):
+        nonlocal used_allocs
+        e = c._node
+        if isinstance(e, LoopIR.Read):
+            if not isinstance(e.idx[dim], LoopIR.Const):
+                raise SchedulingError(
+                    f"Expected a constant buffer access, got {e.idx[dim]} at {dim}'th dimension. Try unrolling the loop. "
+                )
+
+            used_allocs.add(e.idx[dim].val)
+
+            sym = [e.idx[dim].val]
+            new_idx = e.idx.copy()
+            del new_idx[dim]
+
+            return {"name": sym, "idx": new_idx}
+        elif isinstance(e, LoopIR.WindowExpr):
+            if not isinstance(e.idx[dim], LoopIR.Point):
+                raise SchedulingError(
+                    f"Cannot unroll a buffer at a dimension used as a window."
+                )
+
+            pt = e.idx[dim].pt
+            if not isinstance(pt, LoopIR.Const):
+                raise SchedulingError(
+                    f"Expected a constant buffer access, got {pt} at {dim}'th dimension. Try unrolling the loop. "
+                )
+
+            used_allocs.add(pt.val)
+            sym = buf_syms[pt.val]
+            new_access = e.idx.copy()
+            del new_access[dim]
+
+            return {"name": sym, "idx": new_access}
+
+    def mk_write(c):
+        s = c._node
+        if not isinstance(s.idx[dim], LoopIR.Const):
+            raise SchedulingError(
+                f"Expected a constant buffer access, got {s.idx[dim]} at {dim}'th dimension. Try unrolling the loop. "
+            )
+
+        nonlocal used_allocs
+        used_allocs.add(s.idx[dim].val)
+
+        sym = buf_syms[s.idx[dim].val]
+        new_idx = s.idx.copy()
+        del new_idx[dim]
+
+        return {"name": sym, "idx": new_idx}
+
+    ir, fwd = alloc_cursor.get_root(), lambda x: x
+    for c in get_rest_of_block(alloc_cursor):
+        ir, fwd = _replace_pats(ir, fwd, c, f"{alloc_stmt.name}[_]", mk_read)
+        ir, fwd = _replace_pats_stmts(ir, fwd, c, f"{alloc_stmt.name} += _", mk_write)
+        ir, fwd = _replace_pats_stmts(ir, fwd, c, f"{alloc_stmt.name} = _", mk_write)
+
+    new_shape = alloc_stmt.type.shape().copy()
+    del new_shape[dim]
+    new_type = LoopIR.Tensor(new_shape, False, alloc_stmt.type.basetype())
+
+    new_allocs = []
+    for itr in used_allocs:
+        alloc = LoopIR.Alloc(
+            buf_syms[itr],
+            new_type,
+            alloc_stmt.mem,
+            None,
+            alloc_stmt.srcinfo,
+        )
+        new_allocs.append(alloc)
+
+    ir, fwd_repl = fwd(alloc_cursor)._replace(new_allocs)
+    fwd = _compose(fwd_repl, fwd)
+
+    return ir, fwd
+
+
 # --------------------------------------------------------------------------- #
 # --------------------------------------------------------------------------- #
 # The Passes to export
@@ -3652,6 +3750,7 @@ __all__ = [
     "DoBindConfig",
     "DoConfigWrite",
     "DoDeleteConfig",
+    "DoUnrollBuffer",
     ### END Scheduling Ops with Cursor Forwarding ###
     "DoPartialEval",
     "DoExtractMethod",
diff --git a/src/exo/platforms/rvv.py b/src/exo/platforms/rvv.py
index a464dc3..72bb883 100644
--- a/src/exo/platforms/rvv.py
+++ b/src/exo/platforms/rvv.py
@@ -30,7 +30,6 @@ class RVV(Memory):
         if not shape:
             raise MemGenError(f"{srcinfo}: RVV vectors are not scalar values")
 
-
         vec_types = {
             "float": (4, "vfloat32m1_t")
         }  # , "double": (2, "float64x2_t"), "_Float16" : (8, "float16x8_t")}
@@ -72,7 +71,6 @@ class RVV(Memory):
         return f"{baseptr}{idxs}"
 
 
-
 # --------------------------------------------------------------------------- #
 #   f32 RVV intrinsics
 # --------------------------------------------------------------------------- #
@@ -104,6 +102,7 @@ def rvv_vst_4xf32(dst: [f32][4] @ DRAM, src: [f32][4] @ RVV, vl: size):
     for i in seq(0, vl):
         dst[i] = src[i]
 
+
 @instr("{dst_data} = __riscv_vfmv_v_f_f32m1({src_data},{vl});")
 def rvv_broadcast_4xf32(dst: [f32][4] @ RVV, src: [f32][1] @ DRAM, vl: size):
     assert stride(dst, 0) == 1
@@ -147,6 +146,7 @@ def rvv_vfmacc_4xf32_4xf32(
     for i in seq(0, vl):
         dst[i] += lhs[i] * rhs[i]
 
+
 @instr("{dst_data} = __riscv_vfmacc_vf_f32m1{dst_data}, {rhs_data}, {lhs_data},{vl});")
 def rvv_vfmacc_4xf32_1xf32(
     dst: [f32][4] @ RVV, lhs: [f32][4] @ RVV, rhs: [f32][1] @ DRAM, vl: size
@@ -173,4 +173,3 @@ def rvv_vfmacc_1xf32_4xf32(
 
     for i in seq(0, vl):
         dst[i] += lhs[0] * rhs[i]
-
diff --git a/src/exo/query_asts.py b/src/exo/query_asts.py
index 5af4043..1b69e27 100644
--- a/src/exo/query_asts.py
+++ b/src/exo/query_asts.py
@@ -97,6 +97,7 @@ class Stmt(QueryAST):
 class R(Type):
     pass
 
+
 @_dataclass
 class f16(Type):
     pass
diff --git a/src/exo/stdlib/scheduling.py b/src/exo/stdlib/scheduling.py
index b596725..2398bcc 100644
--- a/src/exo/stdlib/scheduling.py
+++ b/src/exo/stdlib/scheduling.py
@@ -56,6 +56,7 @@ from ..API_scheduling import (
     inline_window,
     stage_window,
     stage_mem,
+    unroll_buffer,
     #
     # loop rewriting
     divide_loop,


Test patch:
diff --git a/tests/golden/test_cursors/test_unroll_buffer_forwarding.txt b/tests/golden/test_cursors/test_unroll_buffer_forwarding.txt
new file mode 100644
index 0000000..83224fd
--- /dev/null
+++ b/tests/golden/test_cursors/test_unroll_buffer_forwarding.txt
@@ -0,0 +1,11 @@
+def foo():
+    ...
+    src_0 = 1.0  # <-- NODE
+    ...
+def foo():
+    ...
+    src_1 = 1.0  # <-- NODE
+def foo():
+    ...
+    src_0 = 1.0  # <-- NODE
+    ...
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_unroll_buffer.txt b/tests/golden/test_schedules/test_unroll_buffer.txt
new file mode 100644
index 0000000..7d48628
--- /dev/null
+++ b/tests/golden/test_schedules/test_unroll_buffer.txt
@@ -0,0 +1,16 @@
+def bar(n: size, A: i8[n] @ DRAM):
+    for i in seq(0, n):
+        for j in seq(0, n):
+            tmp_a_0: i8[2] @ DRAM
+            tmp_a_1: i8[2] @ DRAM
+            tmp_a_0[1] = A[i]
+            tmp_a_0[1] = A[i]
+            tmp_a_1[0] = A[i]
+def bar(n: size, A: i8[n] @ DRAM):
+    for i in seq(0, n):
+        for j in seq(0, n):
+            tmp_a_0: i8[5] @ DRAM
+            tmp_a_1: i8[5] @ DRAM
+            tmp_a_1[0] = A[i]
+            tmp_a_1[0] = A[i]
+            tmp_a_0[1] = A[i]
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_unroll_buffer1.txt b/tests/golden/test_schedules/test_unroll_buffer1.txt
new file mode 100644
index 0000000..a718a93
--- /dev/null
+++ b/tests/golden/test_schedules/test_unroll_buffer1.txt
@@ -0,0 +1,8 @@
+def bar(n: size, A: i8[n] @ DRAM):
+    assert n > 10
+    for i in seq(0, n - 4):
+        for j in seq(0, n):
+            tmp_a_0: i8[4, 2] @ DRAM
+            tmp_a_1: i8[4, 2] @ DRAM
+            foo(tmp_a_0[0, 0:2], A[i:i + 2])
+            foo(tmp_a_1[0, 0:2], A[i + 2:i + 4])
\ No newline at end of file
diff --git a/tests/test_cursors.py b/tests/test_cursors.py
index b68dcfe..b253cba 100644
--- a/tests/test_cursors.py
+++ b/tests/test_cursors.py
@@ -257,15 +257,26 @@ def test_vectorize_forwarding(golden):
     scal4 = lift_alloc(scal3, "alphaReg")
     scal5 = fission(scal4, scal4.find("alphaReg[_] = _").after())
 
-    scal1.forward(stmt)
-    scal2.forward(stmt)
-    scal3.forward(stmt)
-    scal4.forward(stmt)
-    scal5.forward(stmt)
-
     assert str(scal5.forward(stmt)) == golden
 
 
+def test_unroll_buffer_forwarding(golden):
+    @proc
+    def foo():
+        src: i32[2]
+        src[0] = 1.0
+        src[1] = 1.0
+
+    assn1 = foo.find("src[_] = _")
+    const1 = assn1.rhs()
+    assn2 = foo.find("src[_] = _ #1")
+    foo1 = unroll_buffer(foo, "src: _", 0)
+
+    tests = [foo1.forward(assn1), foo1.forward(assn2), foo1.forward(const1).parent()]
+
+    assert "\n".join([str(test) for test in tests]) == golden
+
+
 def test_arg_cursor(golden):
     @proc
     def scal(n: size, alpha: R, x: [R][n, n]):
diff --git a/tests/test_neon.py b/tests/test_neon.py
index bbe0dd6..92840ce 100644
--- a/tests/test_neon.py
+++ b/tests/test_neon.py
@@ -131,10 +131,12 @@ def test_neon_vfmla():
 
     return simplify(simple_neon_vfmla)
 
+
 @pytest.mark.isa("neon")
 def test_gen_neon_vfmla(golden, test_neon_vfmla):
     assert str(test_neon_vfmla) == golden
 
+
 @pytest.fixture
 def test_neon_vfmla_f16():
     """
diff --git a/tests/test_rvv.py b/tests/test_rvv.py
index efd9ff0..0a2b7c9 100644
--- a/tests/test_rvv.py
+++ b/tests/test_rvv.py
@@ -63,6 +63,7 @@ def test_rvv_vfmul():
 
     return simplify(simple_rvv_vfmul)
 
+
 @pytest.mark.isa("rvv")
 def test_gen_rvv_vfmul(golden, test_rvv_vfmul):
     assert str(test_rvv_vfmul) == golden
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 1b9f393..ba91baf 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -2715,3 +2715,100 @@ def test_extract_subproc(golden):
         foo, "fooooo", "for i in _:_", order={"x": 1, "y": 0, "j": 2}
     )
     assert (str(foo) + "\n" + str(new)) == golden
+
+
+def test_unroll_buffer(golden):
+    @proc
+    def bar(n: size, A: i8[n]):
+        for i in seq(0, n):
+            for j in seq(0, n):
+                tmp_a: i8[5, 2]
+                tmp_a[0, 1] = A[i]
+                tmp_a[0, 1] = A[i]
+                tmp_a[1, 0] = A[i]
+
+    bar0 = unroll_buffer(bar, "tmp_a : _", 0)
+    bar1 = unroll_buffer(bar, "tmp_a : _", 1)
+    assert str(bar0) + "\n" + str(bar1) == golden
+
+
+def test_unroll_buffer1(golden):
+    @proc
+    def foo(src: i8[2], dst: i8[2]):
+        src[0] = dst[0]
+        src[1] = dst[1]
+
+    @proc
+    def bar(n: size, A: i8[n]):
+        assert n > 10
+        for i in seq(0, n - 4):
+            for j in seq(0, n):
+                tmp_a: i8[4, 2, 2]
+                foo(tmp_a[0, 0, :], A[i : i + 2])
+                foo(tmp_a[0, 1, :], A[i + 2 : i + 4])
+
+    bar = unroll_buffer(bar, "tmp_a : _", 1)
+    assert str(bar) == golden
+
+
+def test_unroll_buffer2():
+    @proc
+    def bar(n: size, A: i8[n]):
+        tmp_a: i8[10]
+        for i in seq(0, n):
+            for j in seq(0, 10):
+                tmp_a[j] = A[i]
+        tmp_b: i8[10]
+        for j in seq(0, 10):
+            tmp_b[j] = tmp_a[j]
+
+    with pytest.raises(
+        SchedulingError,
+        match="Expected a constant buffer access",
+    ):
+        bar = unroll_buffer(bar, "tmp_a : _", 0)
+
+
+def test_unroll_buffer3():
+    @proc
+    def bar(n: size, A: i8[n]):
+        tmp_a: i8[n]
+        for i in seq(0, n):
+            tmp_a[i] = A[i]
+
+    with pytest.raises(
+        SchedulingError,
+        match="Expected a constant buffer dimension",
+    ):
+        bar = unroll_buffer(bar, "tmp_a : _", 0)
+
+
+def test_unroll_buffer4():
+    @proc
+    def bar(n: size, A: i8[n]):
+        tmp_a: i8
+        for i in seq(0, n):
+            tmp_a = A[i]
+
+    with pytest.raises(
+        SchedulingError,
+        match="Cannot unroll a scalar buffer",
+    ):
+        bar = unroll_buffer(bar, "tmp_a : _", 0)
+
+
+def test_unroll_buffer5():
+    @proc
+    def foo(B: i8[2]):
+        B[0] = 0.0
+
+    @proc
+    def bar(n: size, A: i8[n]):
+        tmp_a: i8[10]
+        foo(tmp_a[0:2])
+
+    with pytest.raises(
+        SchedulingError,
+        match="Cannot unroll a buffer at a dimension used as a window",
+    ):
+        bar = unroll_buffer(bar, "tmp_a : _", 0)