[PassConfig] Introduce PassConfig TL_STORAGE_REWRITE_DETECT_INPLACE
#1089
Conversation
- Add kStorageRewriteDetectInplace constant and register the flag with PassContext so C++ code no longer hard-codes the key. - Wire StorageRewrite to include TileLang builtin constants and honor the new config toggle when deciding inplace reuse. - Document the flag across Python surfaces (PassConfigKey, JIT/autotuner docs) with usage guidance and simplified IR examples.
|
👋 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! 🚀 |
|
relevant to issue #994 . |
WalkthroughA new configurable flag TL_STORAGE_REWRITE_DETECT_INPLACE was added across C++ and Python layers. StorageRewrite now reads this PassContext option (default false) and forwards the detect_inplace value to StoragePlanRewriter.Rewrite; tests and docs updated accordingly. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Pass as StorageRewritePass
participant PC as PassContext
participant Rewriter as StoragePlanRewriter
User->>Pass: run StorageRewrite pass
Pass->>PC: Get(kStorageRewriteDetectInplace)
alt option set
PC-->>Pass: detect_inplace = user_value
else default
PC-->>Pass: detect_inplace = false
end
Pass->>Rewriter: Rewrite(plan, detect_inplace)
Rewriter-->>Pass: rewritten_plan
Pass-->>User: pass complete
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
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: 0
🧹 Nitpick comments (1)
testing/python/components/test_storage_rewrite_detect_inplace.py (1)
7-20: Consider renaming function for clarity.The function name
buggy_kernelis somewhat misleading since the kernel isn't actually buggy—it's demonstrating storage rewrite behavior. Consider using a more descriptive name liketest_kernelorstorage_rewrite_test_kernel.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
src/op/builtin.cc(1 hunks)src/op/builtin.h(1 hunks)src/transform/storage_rewrite.cc(3 hunks)testing/python/components/test_storage_rewrite_detect_inplace.py(1 hunks)tilelang/autotuner/param.py(1 hunks)tilelang/jit/__init__.py(1 hunks)tilelang/jit/kernel.py(1 hunks)tilelang/transform/pass_config.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
testing/python/components/test_storage_rewrite_detect_inplace.py (3)
tilelang/jit/__init__.py (1)
jit(237-310)tilelang/language/symbolics.py (1)
symbolic(26-28)tilelang/transform/pass_config.py (1)
PassConfigKey(6-144)
🔇 Additional comments (11)
tilelang/autotuner/param.py (1)
40-40: LGTM! Documentation updated to reference centralized config.The change correctly redirects users to
tilelang.PassConfigKeyfor supported options, reducing duplication and improving maintainability.tilelang/jit/__init__.py (1)
62-62: LGTM! Documentation reference updated.Correctly directs users to
tilelang.transform.PassConfigKeyfor supported options, maintaining consistency with the broader PR changes.tilelang/jit/kernel.py (1)
74-74: LGTM! Docstring updated consistently.The change aligns with the PR's goal of centralizing pass config documentation in
tilelang.PassConfigKey.src/op/builtin.cc (1)
36-36: LGTM! Pass config option registered correctly.The new
kStorageRewriteDetectInplaceboolean option is registered following the established pattern and matches the constant defined insrc/op/builtin.h.tilelang/transform/pass_config.py (1)
72-110: LGTM! Comprehensive config key documentation.The new
TL_STORAGE_REWRITE_DETECT_INPLACEenum member is well-documented with:
- Clear explanation of default behavior (False)
- Concrete IR examples showing aliasing differences
- Proper usage example with PassContext
The string value correctly matches the C++ constant
kStorageRewriteDetectInplacedefined insrc/op/builtin.h.src/op/builtin.h (1)
51-52: LGTM! Config constant defined correctly.The
kStorageRewriteDetectInplaceconstant is properly defined with the value"tl.storage_rewrite_detect_inplace", matching the PythonPassConfigKey.TL_STORAGE_REWRITE_DETECT_INPLACEenum value.testing/python/components/test_storage_rewrite_detect_inplace.py (1)
52-57: Test validates inplace detection correctly.The test appropriately verifies that:
- Without the flag:
"read = (read * 2);"does not appear (separate allocations)- With the flag:
"read = (read * 2);"appears (inplace reuse)The string-based assertion is acceptable for this code generation test, though it's somewhat fragile to code generation changes.
src/transform/storage_rewrite.cc (4)
41-41: Include added for new config constant.The
#include "../op/builtin.h"is necessary to accesskStorageRewriteDetectInplaceused at line 1919.
470-472: Rewrite signature updated to accept inplace flag.The
detect_inplaceparameter is now passed toRewrite()instead of being hardcoded. This enables runtime configuration via PassContext.
1918-1919: Default value conservative and correct.The config defaults to
Bool(false), which keeps distinct temporaries and avoids implicit aliasing. This is a safe default that maintains backward compatibility for users who don't explicitly enable the flag.
1945-1947: Config flag properly propagated to rewriter.The
detect_inplaceflag read from PassContext is correctly passed toStoragePlanRewriter().Rewrite(), enabling user control over inplace buffer detection.
1 participant
This pull request adds a new configuration option to control inplace buffer detection in the StorageRewrite pass, improves documentation for pass configuration options, and introduces a test to verify the feature. The main change is the introduction of the
tl.storage_rewrite_detect_inplaceoption, which allows users to enable or disable aggressive buffer reuse in kernel compilation.New StorageRewrite configuration option:
kStorageRewriteDetectInplaceoption to TVM's pass config registry (src/op/builtin.cc,src/op/builtin.h). [1] [2]src/transform/storage_rewrite.cc). [1] [2]TL_STORAGE_REWRITE_DETECT_INPLACEtoPassConfigKeyenum with a detailed docstring explaining its behavior and usage (tilelang/transform/pass_config.py).Documentation improvements:
tilelang/autotuner/param.py,tilelang/jit/__init__.py, andtilelang/jit/kernel.pyto refer users toPassConfigKeyfor supported options, rather than listing them inline. [1] [2] [3]Testing:
test_storage_rewrite_detect_inplace.py, to verify that enabling the inplace detection option changes kernel code generation as expected.Summary by CodeRabbit
New Features
Tests
Documentation