Skip to content

[Bugfix][Refactor] Always disable light storage reuse #1691

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 3 commits into from
Jan 19, 2026
Merged

[Bugfix][Refactor] Always disable light storage reuse #1691

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b7ba3a7
fix
LeiWang1999 Jan 19, 2026
4d1d7d7
lint fix
LeiWang1999 Jan 19, 2026
b2cb41c
Remove print statement in test_tilelang_issue_1678.py
LeiWang1999 Jan 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 7 additions & 13 deletions src/transform/storage_rewrite.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1109,6 +1109,9 @@ class StoragePlanRewriter : public StmtExprMutator {
// when not divided, no reuse, eg, float4 vs float3
if (e->bits_offset % op_elem_bits != 0)
continue;
// must check element type to avoid type mismatch in codegen
if (e->elem_type != op->dtype.element_of())
continue;
if (reuse_require_exact_matched_dtype && e->elem_type != op->dtype) {
continue;
}
Expand Down Expand Up @@ -1962,21 +1965,12 @@ Pass StorageRewrite() {
ctx->GetConfig<Bool>(kStorageRewriteDetectInplace, Bool(false)).value();
bool enable_reuse = true;
bool reuse_require_exact_matched_dtype = false;
bool merge_static_smem =
ctx->GetConfig<Bool>("tir.merge_static_smem", Bool(false)).value();
AllocateCollector collector;
collector(f->body);
bool has_dynamic = collector.dyn_shmem_allocs_.size() > 1;
if (has_dynamic || merge_static_smem) {
// For IRModule utilizing dynamic shared memory, reuse is not enabled
// Because dynamic doesn't require maintaining the readability and
// it benefits from a more optimized allocation strategy through the
// Pass `MergeSharedMemoryAllocations`.
// When `merge_static_smem` is true, we will reuse and merge shared
// memory in a dedicated pass `MergeSharedMemoryAllocations`.
// And so we don't enable reuse in this pass.
enable_reuse = false;
}
// Always disable reuse currently, for shared memory reuse we depend on
// MergeSharedMemoryAllocations pass, for register reuse we depend on nvcc
// or other compiler its self.
enable_reuse = false;

Optional<Target> target = f->GetAttr<Target>("target");
if (target.defined() && (target.value()->kind->name == "vulkan" ||
Expand Down
27 changes: 27 additions & 0 deletions testing/python/issue/test_tilelang_issue_1678.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# ruff: noqa
import tilelang
import tilelang.testing
import tilelang.language as T


def test_issue_1678():
@tilelang.jit
def qwq():
@T.prim_func
def qwq_kernel():
with T.Kernel(4096, 1, threads=1) as (pid_y, pid_x):
i = T.alloc_var("int32")
i = 1
tmp_row = T.alloc_local((4,), "float32")
amax_local = T.alloc_var("float32")
j = 0
amax_local = T.max(amax_local, tmp_row[j])

return qwq_kernel

kernel = qwq()
print(kernel.get_kernel_source())

@coderabbitai coderabbitai Bot Jan 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

# First, find the test file and verify it exists
find . -name "test_tilelang_issue_1678.py" -type f

Repository: tile-ai/tilelang

Length of output: 110

🏁 Script executed:

# Also list test files in the issue directory to understand the structure
fd "test_.*\.py" testing/python/issue --max-results 20

Repository: tile-ai/tilelang

Length of output: 941

🏁 Script executed:

# Search for similar tests that validate kernel source
rg "get_kernel_source" -A 5 -B 2 testing/python --type py | head -100

Repository: tile-ai/tilelang

Length of output: 8596

🏁 Script executed:

cat -n testing/python/issue/test_tilelang_issue_1678.py

Repository: tile-ai/tilelang

Length of output: 896

🏁 Script executed:

# Check a couple of similar issue tests to see their patterns
rg "def test_issue" -A 20 testing/python/issue/test_tilelang_issue_1001.py | head -60

Repository: tile-ai/tilelang

Length of output: 42

🏁 Script executed:

# Another example
rg "def test_issue" -A 20 testing/python/issue/test_tilelang_issue_1026.py | head -60

Repository: tile-ai/tilelang

Length of output: 233

Add an assertion to validate kernel source generation.

The test currently only prints the kernel source without any assertions. Similar issue tests in the codebase (e.g., test_tilelang_issue_1026.py) validate behavior by asserting on generated source content. Add an assertion to ensure the kernel source is generated correctly and prevent regressions:

     kernel = qwq()
-    print(kernel.get_kernel_source())
+    src = kernel.get_kernel_source()
+    assert src and "qwq_kernel" in src
+    print(src)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def test_issue_1678():
@tilelang.jit
def qwq():
@T.prim_func
def qwq_kernel():
with T.Kernel(4096, 1, threads=1) as (pid_y, pid_x):
i = T.alloc_var("int32")
i = 1
tmp_row = T.alloc_local((4,), "float32")
amax_local = T.alloc_var("float32")
j = 0
amax_local = T.max(amax_local, tmp_row[j])
return qwq_kernel
kernel = qwq()
print(kernel.get_kernel_source())
def test_issue_1678():
`@tilelang.jit`
def qwq():
`@T.prim_func`
def qwq_kernel():
with T.Kernel(4096, 1, threads=1) as (pid_y, pid_x):
i = T.alloc_var("int32")
i = 1
tmp_row = T.alloc_local((4,), "float32")
amax_local = T.alloc_var("float32")
j = 0
amax_local = T.max(amax_local, tmp_row[j])
return qwq_kernel
kernel = qwq()
src = kernel.get_kernel_source()
assert src and "qwq_kernel" in src
print(src)
🤖 Prompt for AI Agents 
In `@testing/python/issue/test_tilelang_issue_1678.py` around lines 7 - 23,
Replace the sole print in test_issue_1678 with an assertion that validates the
generated kernel source: call kernel = qwq(); src = kernel.get_kernel_source()
and assert that src is non-empty and contains relevant identifiers (e.g.,
"amax_local" and "tmp_row" or "max") to ensure the qwq/qwq_kernel JIT produced
the expected code rather than just printing it.



if __name__ == "__main__":
tilelang.testing.main()
Loading