Add unit tests for T.assume #1341
Conversation
|
Caution Review failedThe pull request is closed. WalkthroughAdds a new TileLang test module that JIT-compiles small kernels using Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (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 |
|
👋 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! 🚀 |
![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 (3)
testing/python/language/test_tilelang_language_assume.py (3)
13-20: Rename ambiguouslparameter to improve readabilityThe single-letter
lis easily confused with1and is flagged by ruff (E741). Using a clearer name likeleftorstartmakes the loop and assumes more self‑documenting.- def main( - A: T.Tensor((N, ), "float32"), - l : T.int32, - r : T.int32 - ): + def main( + A: T.Tensor((N,), "float32"), + left: T.int32, + r: T.int32, + ): with T.Kernel(1, threads=32) as _: - for i in T.serial(r - l + 1): - T.assume(l + i >= 0 and l + i < N) - A[l + i] = 0 + for i in T.serial(r - left + 1): + T.assume(left + i >= 0 and left + i < N) + A[left + i] = 0
28-28: Prefernot infor membership assertions (ruff E713)The current pattern
not ("if (" in source)is correct but non‑idiomatic and triggers E713. Usingnot inis clearer and matches our lint rules.- assert not ("if (" in source) + assert "if (" not in source @@ - assert ("float4" in source) and not ("if (" in source) + assert ("float4" in source) and "if (" not in source @@ - assert not ("if (" in source) + assert "if (" not in sourceIf you have ruff enabled on this repo, re-running it after this change should clear the E713 warnings. Please verify locally.
Also applies to: 58-58, 89-89
26-26: Consider removing or gatingprint(source)in testsPrinting the generated source in every test run can make CI logs noisy. If you don’t rely on this output, consider removing these prints or gating them behind a debug flag/env var.
Also applies to: 56-56, 87-87
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
testing/python/language/test_tilelang_language_assume.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/language/test_tilelang_language_assume.py
🧬 Code graph analysis (1)
testing/python/language/test_tilelang_language_assume.py (3)
tilelang/language/symbolics.py (1)
dynamic(10-21)tilelang/language/kernel.py (1)
threads(214-218)tilelang/language/v2/builder.py (1)
source(561-562)
🪛 Ruff (0.14.5)
testing/python/language/test_tilelang_language_assume.py
14-14: Ambiguous variable name: l
(E741)
28-28: Test for membership should be not in
Convert to not in
(E713)
58-58: Test for membership should be not in
Convert to not in
(E713)
89-89: Test for membership should be not in
Convert to not in
(E713)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
🔇 Additional comments (1)
testing/python/language/test_tilelang_language_assume.py (1)
6-90: Solid targeted coverage ofT.assumebehaviorsThe three tests nicely exercise boundary-check elimination, vectorization, and more complex indexing flows by asserting against generated kernel source. The structure with small JIT kernels returning a
@T.prim_funcmain is consistent and easy to extend.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
testing/python/language/test_tilelang_language_assume.py (2)
22-24: Consider removing debug print or making it conditional.The
print(source)statement will clutter test output. Consider removing it or making it conditional on an environment variable for debugging purposes.
27-55: Consider more robust source verification.The assertion pattern checking for
"if (" not in sourceis somewhat fragile—it will fail if any if-statement appears in the generated source, even if unrelated to boundary checks. Consider:
- Using more specific patterns (e.g., checking for absence of specific boundary check conditions)
- Parsing the generated source more robustly
- Adding comments explaining the expected behavior
Also, the parentheses around the string literals in the assertion are unnecessary:
assert "float4" in source and "if (" not in sourcewould be clearer.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
testing/python/language/test_tilelang_language_assume.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/language/test_tilelang_language_assume.py
🧬 Code graph analysis (1)
testing/python/language/test_tilelang_language_assume.py (4)
tilelang/language/symbolics.py (1)
dynamic(10-21)tilelang/language/proxy.py (1)
Tensor(252-253)tilelang/language/kernel.py (1)
threads(214-218)tilelang/language/v2/builder.py (1)
source(561-562)
🪛 Ruff (0.14.5)
testing/python/language/test_tilelang_language_assume.py
12-12: Ambiguous variable name: l
(E741)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
🔇 Additional comments (3)
testing/python/language/test_tilelang_language_assume.py (3)
1-4: LGTM!The imports are appropriate for TileLang testing.
57-86: LGTM with same concerns as previous tests.The complex indexing test logic is sound. However, it shares the same concerns noted in previous tests regarding the print statement and assertion fragility.
88-89: LGTM!The test runner follows the standard TileLang testing pattern.
| N = T.dynamic('N') | ||
|
|
||
| @T.prim_func | ||
| def main(A: T.Tensor((N,), "float32"), l : T.int32, r : T.int32): |
There was a problem hiding this comment.
Rename ambiguous variable l to improve readability.
The variable name l (lowercase L) can be easily confused with 1 (digit one) or I (uppercase i). As per coding guidelines.
Apply this diff to use a clearer name:
- def main(A: T.Tensor((N,), "float32"), l : T.int32, r : T.int32):
+ def main(A: T.Tensor((N,), "float32"), left : T.int32, right : T.int32):
with T.Kernel(1, threads=32) as _:
- for i in T.serial(r - l + 1):
- T.assume(l + i >= 0 and l + i < N)
- A[l + i] = 0
+ for i in T.serial(right - left + 1):
+ T.assume(left + i >= 0 and left + i < N)
+ A[left + i] = 0Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 Ruff (0.14.5)
12-12: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
In testing/python/language/test_tilelang_language_assume.py around line 12,
rename the ambiguous parameter `l` to a clearer name (e.g., `left_idx`, `lo`, or
`start`) in the function signature and update any corresponding references
inside the function and tests to match the new name; ensure type annotation
remains T.int32 and run tests to confirm no references were missed.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
testing/python/language/test_tilelang_language_assume.py (1)
12-16: Rename ambiguous parameterlfor readability (Ruff E741).
l(lowercase L) is easy to confuse with1/I, and Ruff already flags it. Consider renaming to something likeleft/start(andrtoright/end) and updating the uses in this kernel and any related tests.
🧹 Nitpick comments (2)
testing/python/language/test_tilelang_language_assume.py (2)
41-41: Avoid magic constant; reusevectorize_sizeforbase_idx.
base_idx = tid * 4duplicates the literal already stored invectorize_size. Using the named constant keeps things in sync if you ever change the vector width:- base_idx = tid * 4 + base_idx = tid * vectorize_size
24-24: Make source-string assertions more robust to codegen changes.All three tests rely on raw substring checks like
"if (" not in source(and"float4" in source). These can become brittle if codegen formatting or type naming changes (e.g.,if(, different spacing, or a different vector type spelling) even when the semantic property you care about is unchanged.If the APIs are available, consider instead:
- Inspecting the underlying IR/module to assert there are no conditional
If/IfThenElsenodes afterT.assumeis applied.- Or, at minimum, using a regex like
r"\bif\s*\("for the string search, and a more semantic check for vector lanes instead of the exact"float4"token.This would make the tests more future-proof while still validating the intended optimization behavior.
Also applies to: 54-54, 85-85
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
testing/python/language/test_tilelang_language_assume.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/language/test_tilelang_language_assume.py
🧬 Code graph analysis (1)
testing/python/language/test_tilelang_language_assume.py (4)
tilelang/language/symbolics.py (1)
dynamic(10-21)tilelang/language/proxy.py (1)
Tensor(252-253)tilelang/language/kernel.py (1)
threads(214-218)tilelang/language/v2/builder.py (1)
source(561-562)
🪛 Ruff (0.14.5)
testing/python/language/test_tilelang_language_assume.py
12-12: Ambiguous variable name: l
(E741)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
🔇 Additional comments (1)
testing/python/language/test_tilelang_language_assume.py (1)
6-87: Overall: tests exercise the keyT.assumebehaviors and look sound.The three tests nicely cover: (1) removal of explicit boundary checks, (2) enabling vectorized access patterns, and (3) interaction with more complex indexing, all via the JIT path and
get_kernel_source(). Aside from minor readability and robustness nits already mentioned, there are no blocking issues from this file.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
testing/python/language/test_tilelang_language_assume.py (1)
13-17: Rename ambiguous parameterl(E741) and keep bounds expressions clear.Ruff flags
las an ambiguous name, and it’s hard to distinguish from1/I. Renaming also clarifies the left/right bound semantics.Consider:
- def main(A: T.Tensor((N,), "float32"), l: T.int32, r: T.int32): + def main(A: T.Tensor((N,), "float32"), left: T.int32, right: T.int32): with T.Kernel(1, threads=32) as _: - for i in T.serial(r - l + 1): - T.assume(l + i >= 0 and l + i < N) - A[l + i] = 0 + for i in T.serial(right - left + 1): + T.assume(left + i >= 0 and left + i < N) + A[left + i] = 0
🧹 Nitpick comments (3)
testing/python/language/test_tilelang_language_assume.py (3)
31-48: Usevectorize_sizeconsistently and avoid hard‑coded “4”.
vectorize_sizeis defined butbase_idxstill uses a literal4, which can drift if you change the vector width later.- base_idx = tid * 4 + base_idx = tid * vectorize_sizeThis keeps the indexing logic tied to the vectorization factor.
23-25: Reduce brittleness/noise fromprint(source)and global"if ("string checks.All three tests print the full kernel source and assert on
"if ("over the entire text; this can clutter test output and make tests sensitive to unrelated formatting or backend changes.Consider:
- Dropping
print(source)(or only printing on assertion failure), and- Narrowing assertions (e.g., helper that checks only the relevant kernel body, or at least splitting checks so failures are easier to interpret).
Also applies to: 54-56, 86-88
73-80: Document intent of complexi_src/j_srcindexing to aid maintainability.The combination of
T.min,T.ceildiv, and chained assumes makes the indexing logic non‑obvious. A brief comment explaining that this is intentionally “gnarly” to stressT.assume’s analysis on complex expressions would help future readers.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
testing/python/language/test_tilelang_language_assume.py(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/language/test_tilelang_language_assume.py
🧬 Code graph analysis (1)
testing/python/language/test_tilelang_language_assume.py (3)
tilelang/language/symbolics.py (1)
dynamic(10-21)tilelang/language/proxy.py (1)
Tensor(252-253)tilelang/language/v2/builder.py (1)
source(561-562)
🪛 Ruff (0.14.5)
testing/python/language/test_tilelang_language_assume.py
13-13: Ambiguous variable name: l
(E741)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
Remove print statement for kernel source in tests.
2 participants
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.