[Feature] Warp specialization schedules and materialization - #2892
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! 🚀 |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe PR adds warp-specialization schedule contracts, annotation forwarding, CUDA schedule materialization, validation tests, and configurable AWS GEMM and FlashAttention examples. ChangesWarp-specialization scheduling
AWS CUDA examples
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
tilelang/tilelang/language/copy_op.py
Line 62 in f17f516
Preserve copy annotations on scalar copies
The new annotations parameter is silently ignored whenever both operands are scalar BufferLoads, because that fast path returns a plain BufferStore before the annotations are normalized or attached. In a scheduled kernel, T.copy(A[i], B[j], annotations={T.WSID: "copy"}) therefore loses its tl.ws_op_id and the materializer later reports the store as an unannotated statement, so one-element copies cannot be scheduled through the advertised copy API unless users know to avoid this path or wrap it manually.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/language/copy_op.py (1)
110-121: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winScalar fast paths discard the new
annotationsparameter. Each of these wrappers returns from a scalar early-exit branch before it assembles the annotation dict, so a caller-suppliedannotationsvalue is silently dropped. This breaks the new warp-specialization flow: an op whosetl.ws_op_idannotation disappears cannot be referenced from aWSScopebody, and the error appears later insideMaterializeWSScheduleinstead of at the call site.
tilelang/language/copy_op.py#L110-L121: incopy, either forwardannotationsthrough theBufferStorescalar path or raise an error whenannotationsis non-empty; apply the same treatment to theasync_copyscalar branch at Lines 223-224.tilelang/language/atomic.py#L84-L96: inatomic_max, forwardannotationsto the scalaratomic_max_elem_opcall or raise an error whenannotationsis non-empty; apply the same treatment to the scalar branches ofatomic_min(Lines 169-181) andatomic_add(Lines 265-279).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/copy_op.py` around lines 110 - 121, The scalar fast paths in copy, async_copy, atomic_max, atomic_min, and atomic_add discard caller-supplied annotations; update each branch to forward annotations to the resulting BufferStore or atomic element operation, or explicitly reject non-empty annotations consistently. Apply the changes in tilelang/language/copy_op.py lines 110-121 and 223-224, and tilelang/language/atomic.py lines 84-96, 169-181, and 265-279, preserving existing scalar behavior when annotations are empty.
🧹 Nitpick comments (8)
src/cuda/transform/materialize_ws_schedule.cc (1)
1682-1709: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse the already computed
minlocal in the emittedFor.Line 1685 computes
PrimExpr min = Substitute(orig->min, ctx.subs), and Line 1706 computes the same substitution again. Theminlocal is then used only forctx.chain. Reuse it so the chain level and the emitted loop bound cannot diverge ifctx.subschanges.♻️ Proposed refactor
- return For(fresh, Substitute(orig->min, ctx.subs), std::move(emit_extent), - orig->kind, + return For(fresh, min, std::move(emit_extent), orig->kind, body.size() == 1 ? body[0] : SeqStmt(std::move(body)), std::nullopt, std::move(ann));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/transform/materialize_ws_schedule.cc` around lines 1682 - 1709, In the loop emission block, update the `For` construction to reuse the already substituted `min` local computed before `ctx.chain.push_back`, instead of calling `Substitute(orig->min, ctx.subs)` again. Keep `ctx.chain` and the emitted loop bound based on this same `min` value.testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py (1)
312-312: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse raw strings for the two
match=patterns that contain regex metacharacters.Ruff reports RUF043 on both lines. Line 312 uses the alternation
|and Line 1350 uses\\s+. Both patterns behave correctly today, but a raw string states the regex intent directly and clears the lint.♻️ Proposed refactor
- with pytest.raises(Exception, match="carries no ws op id|never places it"): + with pytest.raises(Exception, match=r"carries no ws op id|never places it"):- with pytest.raises(Exception, match="one\\s+dimension"): + with pytest.raises(Exception, match=r"one\s+dimension"):Also applies to: 1350-1350
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py` at line 312, Update the pytest.raises match patterns in the relevant tests, including the cases near the existing “carries no ws op id|never places it” assertion and the pattern using “\s+”, to use raw string literals. Preserve the regex patterns and exception behavior while clearing Ruff RUF043.Source: Linters/SAST tools
tilelang/language/ws_schedule.py (2)
64-65: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the constructor/attribute name difference for the warp range.
The constructor takes
warps_loandwarps_hi, but the reflected fields arewarp_loandwarp_hi, so callers readrole.warp_loafter passingwarps_lo=.... The testtest_ws_role_requires_keyword_rangerelies on both spellings. Add one docstring line that states the attribute names, so users do not search forrole.warps_lo.Also consider
int(max_nreg)for symmetry with the other two casts.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/ws_schedule.py` around lines 64 - 65, Add a constructor docstring line near __init__ documenting that the warps_lo and warps_hi parameters are exposed as the warp_lo and warp_hi attributes. Also cast max_nreg to int when passing it to _ffi_api.WSRole, matching the existing casts for the warp bounds.
21-30: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to clear the Ruff RUF022 warning.Ruff reports
__all__is not sorted. Apply isort-style ordering.♻️ Proposed ordering
__all__ = [ - "WSRole", - "WSPipeline", "WSInstr", "WSOpRef", - "WSSync", - "WSSyncKind", - "WSScope", + "WSPipeline", + "WSRole", "WSSchedule", + "WSScope", + "WSSync", + "WSSyncKind", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/ws_schedule.py` around lines 21 - 30, Sort the names in the module-level __all__ declaration using isort-style ordering so Ruff RUF022 passes, while retaining every existing export.Source: Linters/SAST tools
tilelang/language/atomic.py (1)
117-121: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the annotation precedence with the copy APIs, or document the difference.
T.copyandT.copy_clusterintilelang/language/copy_op.pyapply individual arguments only when the key is absent, and their docstrings state that annotation values win. Hereann["memory_order"]overwrites a user-suppliedmemory_orderannotation. The same pattern appears at Lines 202-206 and at Lines 304-307 foruse_tma.Pick one rule for the tile-op APIs. Also add
annotationsto theParameterssections ofatomic_max,atomic_min, andatomic_add; the new parameter is currently undocumented.♻️ Proposed precedence guard (apply to all three functions)
ann = _normalize_annotations(annotations) - if memory_order is not None: + if "memory_order" not in ann and memory_order is not None: ann["memory_order"] = _MEMORY_ORDER_ID_MAP[memory_order]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/atomic.py` around lines 117 - 121, Align the tile-op annotation precedence with T.copy and T.copy_cluster: in atomic_max, atomic_min, and atomic_add, only populate memory_order and use_tma from explicit arguments when the corresponding key is absent from the normalized annotations, preserving user-provided annotation values. Update each function’s Parameters documentation to describe the annotations argument.examples/aws/gemm.py (1)
55-55: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the commented-out mbarrier code.
Lines 55, 220, and 223 keep a disabled
mbarpath. This example is reference documentation, so the leftover lines invite confusion about whether the barrier is required. Delete them, or add a comment that states why the manual barrier is kept for reference.Also applies to: 220-220, 223-223
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/aws/gemm.py` at line 55, Remove the disabled mbarrier remnants around the commented mbar declaration and the related mbar references at the other identified locations in the GEMM example. Keep the example free of inactive manual-barrier code unless an explanatory reference comment is explicitly needed.examples/aws/test_example_aws.py (1)
10-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider skipping the benchmark inside the tests.
Each test calls the example
main(), which runs the correctness check and then twodo_benchcalls withbackend="cupti". The benchmark adds runtime and a CUPTI dependency to CI without improving coverage. Add abenchmark=Falseparameter to bothmain()functions, and pass it here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/aws/test_example_aws.py` around lines 10 - 33, Add a benchmark toggle to both example main() functions used by test_example_gemm/test_example_gemm_ws and test_example_flash_attention/test_example_flash_attention_ws, defaulting to the current benchmark behavior while allowing benchmark=False. Guard the do_bench calls with this parameter, then pass benchmark=False from all four tests so they retain correctness coverage without running benchmarks.examples/aws/flash_attention.py (1)
728-737: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueMake the SDPA inputs contiguous before the benchmark.
Lines 729-731 produce non-contiguous views.
F.scaled_dot_product_attentionaccepts them, but the backend selection can differ for non-contiguous strides, so the reported Torch latency may not reflect the best available kernel. Call.contiguous()on the permuted tensors, and keep the copy outside the timed lambda.♻️ Proposed change
- q4 = q.permute(0, 2, 1, 3) - k4 = k.permute(0, 2, 1, 3) - v4 = v.permute(0, 2, 1, 3) + q4 = q.permute(0, 2, 1, 3).contiguous() + k4 = k.permute(0, 2, 1, 3).contiguous() + v4 = v.permute(0, 2, 1, 3).contiguous()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/aws/flash_attention.py` around lines 728 - 737, Update the q4, k4, and v4 assignments in the SDPA benchmark to call contiguous() after permute(), keeping these copies before the timed torch_latency lambda so only attention execution is measured.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/aws/flash_attention.py`:
- Around line 680-687: Update reference_attention to avoid materializing full
scores and probabilities for all heads and sequence positions at once. Compute
the reference attention per head, and chunk the query dimension further if
needed so intermediate memory remains bounded while preserving the existing
float32 CPU computation and bfloat16 output shape.
In `@examples/aws/gemm.py`:
- Around line 48-49: Add an assertion that block_N is evenly divisible by
store_block_N in both the gemm and gemm_ws entry points, alongside the existing
divisibility checks, so the CLI-provided dimensions guarantee every computed
store slice remains within block_N.
In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py`:
- Around line 123-126: Replace the permissive alternative in the version-index
assertion within the relevant materialize workspace schedule test with one regex
that requires A_shared’s first index to be the acquired phase expression i % 3,
followed by the buffer’s remaining indices. Preserve the existing
allocation-shape assertion and align the pattern’s precision with the exact
version-index check used by test_op_node_loop.
In `@tilelang/cuda/pipeline.py`:
- Around line 95-99: Update the warp-specialization handling around
allow_warp_specialized and MaterializeWSSchedule so explicit
T.annotate_ws_schedule annotations and tl.ws_op_id markers are rejected or
materialized whenever warp specialization is disabled. Ensure non-CUDA, no-TMA,
and tl.disable_warp_specialized paths do not silently ignore these annotations;
either validate and raise before lowering or apply MaterializeWSSchedule through
the same required paths, while preserving the existing
ProducerConsumerWarpSpecialized flow when enabled.
In `@tilelang/language/scan_op.py`:
- Line 90: Propagate the annotations argument from the scan entry points through
the fragment branches, including cumsum_fragment and cummax_fragment, into
_scan_fragment. Ensure _scan_fragment attaches the received annotations to the
emitted scan operation so fragment scans preserve T.WSID and other metadata.
---
Outside diff comments:
In `@tilelang/language/copy_op.py`:
- Around line 110-121: The scalar fast paths in copy, async_copy, atomic_max,
atomic_min, and atomic_add discard caller-supplied annotations; update each
branch to forward annotations to the resulting BufferStore or atomic element
operation, or explicitly reject non-empty annotations consistently. Apply the
changes in tilelang/language/copy_op.py lines 110-121 and 223-224, and
tilelang/language/atomic.py lines 84-96, 169-181, and 265-279, preserving
existing scalar behavior when annotations are empty.
---
Nitpick comments:
In `@examples/aws/flash_attention.py`:
- Around line 728-737: Update the q4, k4, and v4 assignments in the SDPA
benchmark to call contiguous() after permute(), keeping these copies before the
timed torch_latency lambda so only attention execution is measured.
In `@examples/aws/gemm.py`:
- Line 55: Remove the disabled mbarrier remnants around the commented mbar
declaration and the related mbar references at the other identified locations in
the GEMM example. Keep the example free of inactive manual-barrier code unless
an explanatory reference comment is explicitly needed.
In `@examples/aws/test_example_aws.py`:
- Around line 10-33: Add a benchmark toggle to both example main() functions
used by test_example_gemm/test_example_gemm_ws and
test_example_flash_attention/test_example_flash_attention_ws, defaulting to the
current benchmark behavior while allowing benchmark=False. Guard the do_bench
calls with this parameter, then pass benchmark=False from all four tests so they
retain correctness coverage without running benchmarks.
In `@src/cuda/transform/materialize_ws_schedule.cc`:
- Around line 1682-1709: In the loop emission block, update the `For`
construction to reuse the already substituted `min` local computed before
`ctx.chain.push_back`, instead of calling `Substitute(orig->min, ctx.subs)`
again. Keep `ctx.chain` and the emitted loop bound based on this same `min`
value.
In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py`:
- Line 312: Update the pytest.raises match patterns in the relevant tests,
including the cases near the existing “carries no ws op id|never places it”
assertion and the pattern using “\s+”, to use raw string literals. Preserve the
regex patterns and exception behavior while clearing Ruff RUF043.
In `@tilelang/language/atomic.py`:
- Around line 117-121: Align the tile-op annotation precedence with T.copy and
T.copy_cluster: in atomic_max, atomic_min, and atomic_add, only populate
memory_order and use_tma from explicit arguments when the corresponding key is
absent from the normalized annotations, preserving user-provided annotation
values. Update each function’s Parameters documentation to describe the
annotations argument.
In `@tilelang/language/ws_schedule.py`:
- Around line 64-65: Add a constructor docstring line near __init__ documenting
that the warps_lo and warps_hi parameters are exposed as the warp_lo and warp_hi
attributes. Also cast max_nreg to int when passing it to _ffi_api.WSRole,
matching the existing casts for the warp bounds.
- Around line 21-30: Sort the names in the module-level __all__ declaration
using isort-style ordering so Ruff RUF022 passes, while retaining every existing
export.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ca0549fb-e9a6-4ba3-a600-e86177cda306
📒 Files selected for processing (21)
examples/aws/flash_attention.pyexamples/aws/gemm.pyexamples/aws/test_example_aws.pysrc/cuda/transform/materialize_ws_schedule.ccsrc/ir.ccsrc/transform/common/warp_specialize.htesting/python/transform/test_tilelang_transform_materialize_ws_schedule.pytilelang/cuda/pipeline.pytilelang/cuda/transform/__init__.pytilelang/language/annotations.pytilelang/language/atomic.pytilelang/language/common.pytilelang/language/copy_op.pytilelang/language/experimental/gemm_sp_op.pytilelang/language/fill_op.pytilelang/language/gemm_op.pytilelang/language/loop.pytilelang/language/reduce_op.pytilelang/language/scan_op.pytilelang/language/utils.pytilelang/language/ws_schedule.py
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py (1)
800-801: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse unpacking instead of list concatenation.
Ruff reports RUF005 on both lines. Replace the
+concatenation with an unpacked list literal.♻️ Proposed fix
- "Producer": producer + ["sched_next"], - "Consumer": consumer + ["sched_next"], + "Producer": [*producer, "sched_next"], + "Consumer": [*consumer, "sched_next"],🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py` around lines 800 - 801, Update the Producer and Consumer list construction to use unpacking in list literals instead of + concatenation, preserving the existing producer, consumer, and "sched_next" elements while resolving Ruff RUF005.Source: Linters/SAST tools
src/cuda/op/copy_analysis.cc (1)
856-871: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared
prefer_instructiondispatch into one helper.Lines 837-871 duplicate lines 759-793 of
SelectCopyInstForLoweringalmost verbatim. The only difference is theallow_storeargument passed toSelectTmaInst. The conflict messages, theIsAutoAsyncCopyEnabledcheck, and thekSyncfallthrough are identical strings.Two copies of the same message text will diverge. Extract a helper such as
std::optional<CopyInstSelection> SelectPreferredCopyInst(const CopyFacts &facts, bool allow_store)in the anonymous namespace and call it from both sites. The helper returnsstd::nulloptforPreferredCopyInstruction::kAutoso each caller keeps its own default path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/op/copy_analysis.cc` around lines 856 - 871, Extract the duplicated prefer_instruction dispatch from SelectCopyInstForLowering and the surrounding lowering path into an anonymous-namespace helper such as SelectPreferredCopyInst(const CopyFacts&, bool allow_store). Move the shared kCPAsync validation, conflict messages, async availability handling, and kSync selection into the helper, passing allow_store through to SelectTmaInst where needed. Return std::nullopt for PreferredCopyInstruction::kAuto, and have both callers retain their existing default-selection logic.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/cuda/op/copy_analysis.cc`:
- Around line 846-854: Update ClassifyCall to inspect the CopyInstSelection
returned by ClassifyWarpSpecializedProducerCopy, checking sel.supported before
relying on sel.inst. When unsupported, return Unsupported(sel.reason) so
prefer_instruction="tma" and "cp_async" failures are surfaced instead of being
scheduled as OpAtom::kSync; preserve the existing supported selection path.
In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py`:
- Around line 411-413: Scope the shifted-index assertion to the producer branch
by moving the “i - 1” check after the existing producer slice created at line
417. Assert the stage-1 buffer access specifically uses the shifted iteration,
rather than searching the entire script, while preserving the existing loop and
invalid-condition assertions.
- Around line 1028-1032: Strengthen the grouped-statement check in this test by
asserting the two expected A_frag statements individually within the consumer
branch, rather than relying on script.count("A_frag") >= 2. Preserve the
existing checks for the rebound A_shared access and absence of ws_op_id.
---
Nitpick comments:
In `@src/cuda/op/copy_analysis.cc`:
- Around line 856-871: Extract the duplicated prefer_instruction dispatch from
SelectCopyInstForLowering and the surrounding lowering path into an
anonymous-namespace helper such as SelectPreferredCopyInst(const CopyFacts&,
bool allow_store). Move the shared kCPAsync validation, conflict messages, async
availability handling, and kSync selection into the helper, passing allow_store
through to SelectTmaInst where needed. Return std::nullopt for
PreferredCopyInstruction::kAuto, and have both callers retain their existing
default-selection logic.
In `@testing/python/transform/test_tilelang_transform_materialize_ws_schedule.py`:
- Around line 800-801: Update the Producer and Consumer list construction to use
unpacking in list literals instead of + concatenation, preserving the existing
producer, consumer, and "sched_next" elements while resolving Ruff RUF005.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 7399ac8b-3662-48f1-90d8-1a4cbb207dd9
📒 Files selected for processing (6)
examples/aws/flash_attention.pyexamples/aws/gemm.pysrc/cuda/op/copy_analysis.ccsrc/cuda/transform/materialize_ws_schedule.cctesting/python/transform/test_tilelang_transform_materialize_ws_schedule.pytilelang/language/annotations.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tilelang/language/annotations.py
- src/cuda/transform/materialize_ws_schedule.cc
- examples/aws/flash_attention.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/cuda/transform/materialize_ws_schedule.cc (1)
436-449: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject schedule values before narrowing them.
WSScheduleacceptsint64_tfields, but these assignments narrow values tointbefore validation. An out-of-range value can become implementation-defined and pass later checks with a different value. This affectsnum_warps, role ranges,max_nreg, pipeline depth, and sync stage.Use a checked
int64_t-to-intconversion before each assignment. Reject negativemax_nreg. Rejectnum_warpsvalues that overflownum_warps_ * 32.Proposed fix
+int CheckedInt(int64_t value, const char* field) { + ICHECK_GE(value, static_cast<int64_t>(std::numeric_limits<int>::min())) + << "ws_schedule: " << field << " is out of range"; + ICHECK_LE(value, static_cast<int64_t>(std::numeric_limits<int>::max())) + << "ws_schedule: " << field << " is out of range"; + return static_cast<int>(value); +} + - num_warps_ = static_cast<int>(sched->num_warps); + num_warps_ = CheckedInt(sched->num_warps, "num_warps"); + ICHECK_LE(num_warps_, std::numeric_limits<int>::max() / 32); ... - role.warp_lo = static_cast<int>(r->warp_lo); - role.warp_hi = static_cast<int>(r->warp_hi); - role.nreg = static_cast<int>(r->max_nreg); + role.warp_lo = CheckedInt(r->warp_lo, "warp_lo"); + role.warp_hi = CheckedInt(r->warp_hi, "warp_hi"); + role.nreg = CheckedInt(r->max_nreg, "max_nreg"); + ICHECK_GE(role.nreg, 0);Also applies to: 504-505, 530-530
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/cuda/transform/materialize_ws_schedule.cc` around lines 436 - 449, Update the WSSchedule materialization logic around num_warps_, RoleSpec fields, pipeline depth, and sync stage to validate each int64_t value is within int range before narrowing, rejecting invalid values with the existing check mechanism. Explicitly reject negative max_nreg values, and validate num_warps_ before any num_warps_ * 32 calculation can overflow. Apply the checked conversion consistently to role warp_lo/warp_hi, max_nreg, pipeline depth, and sync stage assignments.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@testing/python/language/test_tilelang_language_copy.py`:
- Around line 57-60: Bind annotation assertions to their intended intrinsic
instead of checking values globally: in
testing/python/language/test_tilelang_language_copy.py lines 57-60, verify
“sync” occurs within the T.copy call and “async” within T.async_copy; in
testing/python/language/test_tilelang_language_scan.py lines 63-66, verify
“cumsum” within T.cumsum and “cummax” within T.cummax.
---
Outside diff comments:
In `@src/cuda/transform/materialize_ws_schedule.cc`:
- Around line 436-449: Update the WSSchedule materialization logic around
num_warps_, RoleSpec fields, pipeline depth, and sync stage to validate each
int64_t value is within int range before narrowing, rejecting invalid values
with the existing check mechanism. Explicitly reject negative max_nreg values,
and validate num_warps_ before any num_warps_ * 32 calculation can overflow.
Apply the checked conversion consistently to role warp_lo/warp_hi, max_nreg,
pipeline depth, and sync stage assignments.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: cea39b07-c35a-4ce7-9cb6-d52160ff7f53
📒 Files selected for processing (8)
examples/aws/gemm.pysrc/cuda/transform/materialize_ws_schedule.cctesting/python/language/test_tilelang_language_copy.pytesting/python/language/test_tilelang_language_scan.pytesting/python/transform/test_tilelang_transform_materialize_ws_schedule.pytilelang/language/atomic.pytilelang/language/copy_op.pytilelang/language/scan_op.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tilelang/language/atomic.py
- tilelang/language/copy_op.py
1 participant
This PR adds a new set of annotations for describing arbitrary warp specialization scheme, with the abstraction of pipelines. A pipeline supports 4 kinds of operations:
producer_acquire,producer_commit,consumer_wait, andconsumer_release. Behind the scenes, mbarriers are created to implement the synchronization between warps.This PR introduces a new pass, MaterializeWSSchedule, that transforms an annotated kernel to a warp-specialized kernel.
The set of infrastructure enables further development of fully automatic warp specialization in TileLang. I am working on this.
Summary
MaterializeWSScheduleto lower annotated kernels into warp-specialized CUDA kernels.C++ style / lint notes
docs/developer_guide/cpp_style.md.