Repository: apache/tvm
Pull request: [TIR] Fix lower_warp_memory when there are >1 warp buffers
Fix #5366.

Changes:

1. Fixed the recursion. (Just look at the code)
2. Added a test.

Requesting for a review. @tqchen

Problem statement:
[TIR] lower_warp_memory cannot handle >1 warp buffers
Pass  `lower_warp_memory` cannot handle more than one warp buffers. Buffers except the first one cannot be correctly transformed to warp shuffles.

To reproduce:

```python
import tvm                                                                                                                                  [8/1976]
import topi
import numpy as np

from tvm import te

dtype = "float32"
target = "cuda"
m = 32
A = te.placeholder((m,), name='A', dtype=dtype)
B = te.placeholder((m,), name='B', dtype=dtype)
C = te.compute((m,), lambda i: A[(i + 1) % m] + B[(i + 1) % m], name='C')

cuda_target = tvm.target.create("cuda")
assert m <= cuda_target.thread_warp_size
with cuda_target:
    s = te.create_schedule(C.op)
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")

    AA = s.cache_read(A, "warp", [C])
    BB = s.cache_read(B, "warp", [C])
    xo, xi = s[C].split(C.op.axis[0], nparts=1)
    s[C].bind(xi, tx)
    s[C].bind(xo, bx)
    s[AA].compute_at(s[C], xo)
    s[BB].compute_at(s[C], xo)
    xo, xi = s[AA].split(s[AA].op.axis[0], nparts=1)
    s[AA].bind(xo, bx)
    s[AA].bind(xi, tx)
    xo, xi = s[BB].split(s[BB].op.axis[0], nparts=1)
    s[BB].bind(xo, bx)
    s[BB].bind(xi, tx)

print(tvm.lower(s, [A, B, C], target, simple_mode=True))
compute = tvm.build(s, [A, B, C], target, name="run")
print(compute.imported_modules[0].get_source())
```

I think the problem is `WarpMemoryRewriter::VisitStmt_(const AllocateNode*)` in `lower_warp_memory.cc` doesn't continue the recursion after rewriting the first buffer.

I will fix it.

Hints from issue discussion:

Commit message: fix recursion in lower_warp_memory

Patch to review:
diff --git a/src/tir/transforms/lower_warp_memory.cc b/src/tir/transforms/lower_warp_memory.cc
index 71e7cfaf4832..543d098a654c 100644
--- a/src/tir/transforms/lower_warp_memory.cc
+++ b/src/tir/transforms/lower_warp_memory.cc
@@ -379,7 +379,7 @@ class WarpMemoryRewriter : private StmtMutator {
   Stmt VisitStmt_(const AllocateNode* op) {
     if (warp_buffer_.count(op->buffer_var.get())) {
       WarpAccessRewriter rewriter(warp_size_, &analyzer_);
-      return rewriter.Rewrite(op);
+      return StmtMutator::VisitStmt_(rewriter.Rewrite(op).as<AllocateNode>());
     } else {
       return StmtMutator::VisitStmt_(op);
     }
diff --git a/tests/python/unittest/test_tir_transform_lower_warp_memory.py b/tests/python/unittest/test_tir_transform_lower_warp_memory.py
index a761cf1a95d8..51be480a7cba 100644
--- a/tests/python/unittest/test_tir_transform_lower_warp_memory.py
+++ b/tests/python/unittest/test_tir_transform_lower_warp_memory.py
@@ -132,7 +132,56 @@ def check_cuda(dtype):
     check_cuda("float32")
     check_cuda("float16")
 
+def test_lower_warp_memory_cuda_2_buffers():
+    def check_cuda(dtype):
+        if not tvm.gpu(0).exist or not tvm.runtime.enabled("cuda"):
+            print("skip because cuda is not enabled..")
+            return
+        if dtype == "float16" and not have_fp16(tvm.gpu(0).compute_version):
+            print("Skip because gpu does not have fp16 support")
+            return
+
+        m = 32
+        A = te.placeholder((m,), name='A', dtype=dtype)
+        B = te.placeholder((m,), name='B', dtype=dtype)
+        C = te.compute((m,), lambda i: A[(i + 1) % m] + B[(i + 1) % m], name='C')
+
+        cuda_target = tvm.target.create("cuda")
+        assert m <= cuda_target.thread_warp_size
+        with cuda_target:
+            s = te.create_schedule(C.op)
+            tx = te.thread_axis("threadIdx.x")
+            bx = te.thread_axis("blockIdx.x")
+
+            AA = s.cache_read(A, "warp", [C])
+            BB = s.cache_read(B, "warp", [C])
+            xo, xi = s[C].split(C.op.axis[0], nparts=1)
+            s[C].bind(xi, tx)
+            s[C].bind(xo, bx)
+            s[AA].compute_at(s[C], xo)
+            s[BB].compute_at(s[C], xo)
+            xo, xi = s[AA].split(s[AA].op.axis[0], nparts=1)
+            s[AA].bind(xo, bx)
+            s[AA].bind(xi, tx)
+            xo, xi = s[BB].split(s[BB].op.axis[0], nparts=1)
+            s[BB].bind(xo, bx)
+            s[BB].bind(xi, tx)
+
+            ctx = tvm.gpu(0)
+            func = tvm.build(s, [A, B, C], "cuda")
+            AB_np = np.array(list(range(m)), dtype=dtype)
+            C_np = np.array(list(range(1, m)) + [0], dtype=dtype) * 2
+            A_nd = tvm.nd.array(AB_np, ctx)
+            B_nd = tvm.nd.array(AB_np, ctx)
+            C_nd = tvm.nd.array(np.zeros(C_np.shape, dtype=C_np.dtype), ctx)
+            func(A_nd, B_nd, C_nd)
+            tvm.testing.assert_allclose(C_nd.asnumpy(), C_np, rtol=1e-3)
+
+    check_cuda("float32")
+    check_cuda("float16")
+
 if __name__ == "__main__":
     test_lower_warp_memory_local_scope()
     test_lower_warp_memory_cuda_end_to_end()
     test_lower_warp_memory_cuda_half_a_warp()
+    test_lower_warp_memory_cuda_2_buffers()


Code review comments:
diff_hunk: @@ -379,7 +379,7 @@ class WarpMemoryRewriter : private StmtMutator {
   Stmt VisitStmt_(const AllocateNode* op) {
     if (warp_buffer_.count(op->buffer_var.get())) {
       WarpAccessRewriter rewriter(warp_size_, &analyzer_);
-      return rewriter.Rewrite(op);
+      return StmtMutator::VisitStmt_(rewriter.Rewrite(op).as<AllocateNode>());
path: src/tir/transforms/lower_warp_memory.cc
text: @user1:
Let us consider Post order rewriting pattern instead
```c++
auto ret = StmtMutator::VisitStmt_(op);
op = ret.as<AllocateNode>();
if (warp_buffer_.count(op->buffer_var.get())) {
   // rewrite transformations
}
return ret;
```