[Bugfix][Refactor] Always disable light storage reuse #1691
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Caution Review failedThe pull request is closed. 📝 WalkthroughWalkthroughEnforces element-type checks when considering reusable StorageEntry objects and unconditionally disables storage reuse in the StorageRewrite pass (delegating reuse to downstream passes/compilers). Adds a new TileLang Python test exercising JIT kernel construction and kernel source retrieval. Changes
Sequence Diagram(s)mermaid Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@testing/python/issue/test_tilelang_issue_1678.py`:
- Around line 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.
| 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()) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, find the test file and verify it exists
find . -name "test_tilelang_issue_1678.py" -type fRepository: 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 20Repository: 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 -100Repository: tile-ai/tilelang
Length of output: 8596
🏁 Script executed:
cat -n testing/python/issue/test_tilelang_issue_1678.pyRepository: 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 -60Repository: 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 -60Repository: 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.
| 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.
Remove print statement for kernel source in test.
1 participant
Fix for issue #1678
For shared memory reuse we depend on MergeSharedMemoryAllocations pass, for register reuse we depend on nvcc or other vendors' compilers itself.
Summary by CodeRabbit
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.