mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
6548c05f0c
* [Feature] Batched AllReduce for better T.reduce performance (#1928)
When each thread holds multiple values to reduce, the previous code
emitted one AllReduce call per element inside a loop, causing redundant
synchronization barriers (e.g. 8 values × 6 syncs = 48 syncs).
This change adds batch support to AllReduce so all per-thread values are
reduced in a single butterfly pass sharing barriers (6 syncs total).
- Add `batch_size` and `workspace_stride` template parameters to
`tl::AllReduce` (defaults preserve existing scalar behaviour).
- In ReduceOpNode::Lower, detect when batching is beneficial
(batch_size > 1 and cross-warp reduction) and restructure into three
phases: loop(init + local reduce), batched AllReduce, loop(copy-back).
- Apply same optimisation to FinalizeReducerOpNode::Lower.
- Mirror template changes in HIP AllReduce.
Closes #1928
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Test] Add codegen tests for batched AllReduce
Verify that the batched AllReduce path (batch_size > 1, workspace_stride)
is emitted in the generated CUDA source for reduce_max, reduce_sum, and
reduce_min with cross-warp reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Fix batched AllReduce codegen for HIP/ROCm targets
HIP AllReduce template has no Barrier type parameter, so the batched
path must omit tl::SyncThreadsBarrier for ROCm targets. Adds
TargetIsRocm() checks in both reduce.cc and finalize_reducer.cc.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Fix] Use target-aware warp size for cross-warp reduction threshold
Address CodeRabbit review: ROCm wavefronts are 64-wide, so the cross-warp
threshold should be 64 (not 32) on HIP targets. This avoids unnecessary
workspace allocation and batched codegen for single-wavefront reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Style] Fix ruff format for test assertion
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Support batched AllReduce in CuteDSL Python codegen
The CuteDSL codegen translates C++ AllReduce template calls to Python.
With the new batch_size and workspace_stride template params, the
generated Python calls AllReduce() with 7 args instead of 5. Add
batch_size and workspace_stride params to the Python AllReduce function
and implement batched reduction logic for both standard and Hopper paths.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Feature] Add explicit batch parameter to T.reduce for user-controlled batched AllReduce
Expose a `batch` parameter on T.reduce (and all T.reduce_* wrappers,
default 1 = scalar / current behaviour). When batch > 1 the compiler:
- validates batch divides the per-thread output element count N;
- emits ceil(N/batch) calls to the new AllReduce::run_batch interface,
each sharing a single pair of barriers across all batch elements;
- allocates workspace of size reducing_threads * batch (only for cross-
warp reduce, i.e. reducing_threads > 32).
Key differences from the existing auto-detection approach:
* The user controls when batching is enabled — no silent smem increase.
* workspace_stride = reducing_threads (not block_size), so workspace is
minimal and does not blow past the GPU shared-memory limit.
* run_batch (not run) avoids C++ overload-resolution ambiguity when a
pointer is passed as the first argument.
Files changed:
- src/op/reduce.h: add `int batch{1}` field to ReduceOpNode
- src/op/reduce.cc: read "batch" annotation; validate; emit N/batch
batched AllReduce calls with address_of per-chunk pointers
- src/tl_templates/cuda/reduce.h: rename batch run() -> run_batch(),
fix recursive call
- src/tl_templates/hip/reduce.h: same
- tilelang/language/reduce_op.py: add batch param to reduce() and all
reduce_max/min/sum/absmax/abssum/bitand/bitor/bitxor wrappers
- testing/python/language/test_tilelang_language_reduce.py: update
test_batched_allreduce_codegen to use explicit batch=2 and match
run_batch in generated source
* [Fix] Use run_batch in finalize_reducer batched AllReduce codegen
finalize_reducer.cc also generates batched AllReduce calls (used by
T.alloc_reducer / T.finalize_reducer), but was still emitting ::run
instead of ::run_batch, causing a compiler error because the batched
pointer-based interface is now named run_batch to avoid overload
ambiguity.
* rocm fix
* [Feature] Add explicit batch parameter to T.finalize_reducer
Mirrors the batch= API already available on T.reduce:
- batch=1 (default) keeps the existing scalar AllReduce::run path
- batch>1 emits AllReduce::run_batch, sharing barriers across batch
output elements and reducing barrier count by batch×
C++ changes:
- FinalizeReducerOpNode gains an int batch{1} field (registered via
reflection) and reads it from the 'batch' annotation in the constructor
- Lower() validates batch against the layout output-element count and
uses effective_batch to drive use_batch / codegen
Python changes:
- finalize_reducer(reducer, batch=1) – new optional argument
- batch < 1 raises ValueError (consistent with T.reduce)
Tests: add test_batched_finalize_reducer_codegen covering sum/max/min
with explicit batch values, verifying run_batch is emitted with the
correct template argument.
* [Test] Add comprehensive tests for T.finalize_reducer batch parameter
Covers all cases introduced by the explicit batch= API:
Codegen tests:
- batch>1 emits AllReduce::run_batch with correct template argument
- batch=1 (default) does NOT emit run_batch (scalar path)
Correctness tests:
- batch=1 scalar path: sum/max/min across float32/float16 (PASS)
- batch==block_M full-batch: correct by coincidence (PASS)
- batch>1 partial-batch: xfail(strict=True) – known bug where
run_batch receives buffer->data directly but fragment layout is
non-contiguous per thread; root cause documented in
work/bug-finalize-reducer-batch.md
Error-input tests:
- batch=0 / batch<0 raise ValueError at Python layer
- batch > total output elements raises at compile/lower time
- batch that does not evenly divide output elements raises
* [Test] Rewrite finalize_reducer tests to follow clean outline
Replace the previous over-engineered test suite (4 case lists, 2 split
correctness functions, per-case xfail) with a minimal set organised by
the three categories that matter:
1. Codegen – single parametrised test covers both batch=1 (no run_batch)
and batch>1 (run_batch with correct template arg), reusing one case
list for both codegen and correctness.
2. Correctness – only batch=1 (scalar path), which is currently correct.
batch>1 correctness tests are omitted until the underlying fragment
layout bug is fixed (tracked in work/bug-finalize-reducer-batch.md).
3. Error inputs – batch<1 collapsed into one parametrised test; compile-
time errors (batch>output, batch not divisible) each get one test
using a shared _make_error_kernel helper.
* [Test] Rewrite reduce tests following clean outline
Replace fragmented per-function tests with four parametrized suites:
- test_reduce: op × dtype × src_scope × dst_scope × threads, covers
fragment→fragment, shared→shared, fragment→shared for sum/max/min/
abssum/absmax with float32/float16/int32/int64
- test_reduce_clear: op × src_scope × dst_scope with clear=False,
verifies pre-filled dst is accumulated correctly
- test_reduce_batch_codegen: verify batch>1 emits AllReduce::run_batch
- test_finalize_reducer_codegen/correctness/invalid_batch: batch
parameter feature tests, sharing one FINALIZE_REDUCER_CASES list
Use tilelang.jit lazy style and torch.testing.assert_close directly,
no profiler wrapper.
* [Test] Merge batch codegen+correctness into test_reduce
Add batch column to REDUCE_CASES; batch>1 cases verify both that
run_batch appears in generated source and that numerical output matches
reference, within the same test_reduce parametrization.
Remove the now-redundant test_reduce_batch_codegen.
* [Test] Add more batch>1 reduce cases covering all ops and dtypes
Add batch=2/4/8 cases for sum/max/min/abssum/absmax across float32/
float16/bfloat16, with both fragment and shared src/dst scopes.
Relax tolerance to 1e-1 for float16/bfloat16 to account for larger
rounding error in multi-element reductions.
* [Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug
* Revert "[Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug"
This reverts commit 9d60d6b197.
* fix(HIP): skip __launch_bounds__ for sub-wavefront threads to avoid miscompile of 64-bit shuffle ops
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: kurisu6912 <kurisu6912@users.noreply.github.com>