mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
KellyFrogandchenkf
0bc1913daf
* [CUDA][Reduce] Resolve partial AllReduce participants
* refactor: inline ResolveAllReduceThreadRange into reduce.h and remove allreduce_barriers doc
- Inline ResolveAllReduceThreadRange body from reduce.cc into reduce.h,
add #include <tvm/ir/with_context.h> for With<ConstraintContext> usage.
- Delete src/backend/common/op/reduce.cc (function now in header).
- Remove reduce.cc from CMakeLists.txt TILE_LANG_SRCS.
- Delete docs/compiler_internals/allreduce_barriers.md.
- Remove toctree entry from docs/index.md.
* [CUDA][Reduce] Resolve batched AllReduce participants
* Revert "[CUDA][Reduce] Resolve batched AllReduce participants"
This reverts commit 25ee2101c7.
* [CUDA][Reduce] Fix partial AllReduce warp alignment
---------
Co-authored-by: chenkf <chenkf@clab>
50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import pytest
|
|
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.cuda.target import normalize_cutedsl_target
|
|
|
|
|
|
def _lower_cutedsl_partial_reduce():
|
|
if not tvm.runtime.enabled("cuda"):
|
|
pytest.skip("TileLang CuTeDSL codegen requires TVM built with CUDA support.")
|
|
|
|
build_cutedsl = tvm.ffi.get_global_func("target.build.tilelang_cutedsl_without_compile", allow_missing=True)
|
|
if build_cutedsl is None:
|
|
pytest.skip("TileLang CuTeDSL backend is not enabled in this build.")
|
|
|
|
target = normalize_cutedsl_target({"kind": "cutedsl", "arch": "sm_90"})
|
|
assert target is not None
|
|
|
|
@T.prim_func
|
|
def prog(A: T.Tensor((1, 512), "float32"), B: T.Tensor((1,), "float32")):
|
|
with T.Kernel(1, threads=128):
|
|
x_frag = T.alloc_fragment((1, 512), "float32")
|
|
sum_frag = T.alloc_fragment((1,), "float32")
|
|
T.annotate_layout(
|
|
{
|
|
x_frag: T.Fragment(x_frag.shape, forward_fn=lambda i, j: (j // 8, j % 8)),
|
|
sum_frag: T.Fragment(sum_frag.shape, forward_fn=lambda i, rep: (rep, 0), replicate=64),
|
|
}
|
|
)
|
|
for i, j in T.Parallel(1, 512):
|
|
x_frag[i, j] = A[i, j]
|
|
T.reduce_sum(x_frag, sum_frag, dim=1)
|
|
for i in T.Parallel(1):
|
|
B[i] = sum_frag[i]
|
|
|
|
with target:
|
|
return lower(prog.with_attr("global_symbol", "main"), target=target)
|
|
|
|
|
|
def test_cutedsl_codegen_partial_reduce_named_barrier():
|
|
"""The partial scalar AllReduce uses its exact participant count."""
|
|
artifact = _lower_cutedsl_partial_reduce()
|
|
assert "tl.NamedBarrier(64)" in artifact.kernel_source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|