[CUDA][TMA] Separate atomic-add dtype support from layout encoding - #2846
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! 🚀 |
📝 WalkthroughWalkthroughCUDA TMA layout analysis is centralized for shared-memory encodings and alignment. Atomic-add lowering now validates targets, dtypes, layouts, and descriptor parameters. Bulk-copy lowering uses the same analysis. Atomic-add tests cover additional dtypes, swizzles, targets, and diagnostics. ChangesCUDA TMA layout and atomic add
Estimated code review effort: 4 (Complex) | ~45 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/atomic_add.cc`:
- Around line 449-459: Update the TMA atomic-add validation around
instruction_dim and desc.smem_box in the visible lowering path to ensure the
global_range extent is divisible by the swizzle element width when swizzling is
enabled. Add a clear assertion for this unsupported remainder case, or lower the
descriptor to complete swizzle atoms and explicitly process any remainder
outside TMA; preserve the existing non-swizzled behavior.
In `@testing/python/language/test_tilelang_language_atomic.py`:
- Around line 383-389: Update test_tma_atomic_add_rejects_unsupported_dtype so
the regex interpolates the plain dtype name via str(dtype), matching the C++
diagnostic for TileLang dtype instances while preserving the existing
unsupported-dtype coverage.
🪄 Autofix (Beta)
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: 4ced9527-b0e9-4749-ae05-03e54f0332b3
📒 Files selected for processing (5)
src/cuda/op/atomic_add.ccsrc/cuda/op/copy.ccsrc/cuda/op/tma_layout.ccsrc/cuda/op/tma_layout.htesting/python/language/test_tilelang_language_atomic.py
| desc.swizzle = plan.swizzle_mode.CanonicalOrdinal(); | ||
| RequireTMASmemAlignment(lower_args, shared_tensor, plan.swizzle_mode); | ||
|
|
||
| auto inner_box_dim = as_const_int(desc.smem_box[0]); | ||
| ICHECK(inner_box_dim != nullptr) | ||
| << "inner_box_dim must be a constant integer for TMA atomic add"; | ||
| int instruction_dim = *inner_box_dim; | ||
| if (desc.swizzle == static_cast<int>(CU_TENSOR_MAP_SWIZZLE_64B)) { | ||
| instruction_dim = 64 / shared_tensor->dtype.bytes(); | ||
| } else if (desc.swizzle == static_cast<int>(CU_TENSOR_MAP_SWIZZLE_128B)) { | ||
| instruction_dim = 128 / shared_tensor->dtype.bytes(); | ||
| if (!plan.swizzle_mode.IsNone()) { | ||
| instruction_dim = | ||
| plan.swizzle_mode.ByteWidth() / shared_tensor->dtype.bytes(); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether TMA atomic add constrains the global range extent to the shared tile extent.
set -euo pipefail
echo "=== atomic_add lowering: box vs shared shape ==="
rg -nP -C 6 'smem_box|inner_box_dim|instruction_dim' src/cuda/op/atomic_add.cc
echo "=== range-equality checks in the atomic add op ==="
rg -nP -C 6 'src_range|dst_range' src/op/atomic_add.cc src/cuda/op/atomic_add.cc 2>/dev/null || true
fd -t f 'atomic_add' src --exec rg -nP -C 4 'ICHECK.*(range|extent)' {}
echo "=== tests that atomic-add a sub-tile ==="
rg -nP -C 6 'atomic_add\(.*use_tma' testing/python --glob '*.py'Repository: tile-ai/tilelang
Length of output: 13223
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '=== src/cuda/op/atomic_add.cc relevant sections ==='
sed -n '1,240p' src/cuda/op/atomic_add.cc
printf '\n--- TMA lowering section ---\n'
sed -n '240,570p' src/cuda/op/atomic_add.cc
echo '=== src/op/atomic_add.cc ==='
sed -n '1,180p' src/op/atomic_add.cc
echo '=== TMA atomic add layout helpers ==='
rg -nP -C 8 'MakeTMAAtomicAddSharedLayout|RequireTMASmemAlignment|IsValidTMAReduceAddDtype|SwizzleMode|canonical|canonical ordinal' src/cuda/op/atomic_add.cc src --glob '*.cc' --glob '*.h' --glob '*.hpp' 2>/dev/null || trueRepository: tile-ai/tilelang
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== exact references to src_range/dst_range equality/constraint ==="
rg -n 'src_range|dst_range' src/cuda/op/atomic_add.cc
rg -n 'src_range|dst_range' src/op/atomic_add.cc
rg -n 'dst_range.size.*src_range|src_range.size.*dst_range|src_range.*dst_range|src_range.*==|dst_range.*==' src testing 2>/dev/null | head -80 || true
echo "=== MakeTMAAtomicAddSharedLayout constraints ==="
sed -n '145,177p' src/cuda/op/atomic_add.cc
echo "=== TMA atomic add test files ==="
fd -t f 'atomic_add' tests testing examples 2>/dev/null || true
rg -n -C 4 'tl\.atomic_add.*use_tma|use_tma.*true|atomic_add\(' testing tests examples 2>/dev/null | head -160 || trueRepository: tile-ai/tilelang
Length of output: 19380
🌐 Web query:
CUDA TMA atomic add swizzle descriptor smem_box alignment instruction boundary
💡 Result:
The Tensor Memory Accelerator (TMA) in CUDA requires specific alignment and descriptor configurations to function correctly, particularly when using shared memory (SMEM) swizzling [1][2][3]. 1. Alignment Requirements: TMA destinations in shared memory must be 128-byte aligned [2]. Failure to align the destination buffer to this 128-byte boundary can result in silent data corruption or hardware faults, as there is no runtime check performed by the hardware [2]. 2. SMEM Box and Descriptor Configuration: - The TMA descriptor (a 128-byte opaque structure) encodes the tensor's metadata, including the SMEM box size, strides, and swizzle mode [1][2][3]. - The smem_box defines the shape of the tile as it is laid out in shared memory [1][3]. - When swizzling is applied, the TMA engine rearranges the byte layout in shared memory to avoid bank conflicts [2][3]. The swizzle mode (e.g., SWIZZLE_64B, SWIZZLE_128B) determines the swizzle atom size and pattern [4][5][6]. - The smem_box dimension must be compatible with the selected swizzle mode [4][5]. For example, if a 128B swizzle is used, the instruction dimension in the smem_box calculation often needs to match or be a multiple of the swizzle atom size [4][5]. 3. Instruction Boundaries: The TMA engine operates using atomic "swizzle atoms" [4]. If the inner dimension of the data tile in bytes exceeds the swizzle span (the repeat window of the pattern), the swizzle pattern repeats, which can inadvertently lead to bank conflicts [4]. Consequently, developers must ensure that the box dimensions, swizzle mode, and data types are mathematically consistent so that the TMA hardware can correctly decode the memory layout without violating alignment or bank-access constraints [4][5][7]. In practice, these descriptors are typically generated on the host using the CUDA Driver API (cuTensorMapEncodeTiled) or via libraries like CUTLASS, which automatically manage the mapping between logical tensor coordinates, SMEM swizzle patterns, and TMA hardware requirements [8][2][9][7].
Citations:
- 1: https://docs.nvidia.com/cutlass/4.5.2/media/docs/cpp/cute/0z_tma_tensors.html
- 2: https://nvlabs.github.io/cuda-oxide/advanced/tensor-memory-accelerator.html
- 3: https://mlc.ai/modern-gpu-programming-for-mlsys/chapter_tma/index.html
- 4: https://cudacourseh100.github.io/pages/lesson-4.html
- 5: https://github.com/tile-ai/tilelang-ascend/blob/fe651760/src/op/bulk_copy.cc
- 6: https://github.com/NVIDIA/cutlass/blob/0ce648f5/include/cute/atom/copy_traits_sm90_tma_swizzle.hpp
- 7: https://github.com/NVIDIA/cutlass/blob/d4e16f5d/include/cute/atom/copy_traits_sm90_tma.hpp
- 8: https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/async-copies.html
- 9: [QST]Question about tma load for weights in example92: Is swizzling required? NVIDIA/cutlass#2665
Ensure the atomic-add global range fits the TMA swizzle atom.
instruction_dim is computed from the shared swizzle width, and ICHECK((*inner_box_dim) % instruction_dim == 0) uses desc.smem_box[0], which comes from global_range extents. For TMA global accesses with swizzle, the global tile must also be a multiple of the swizzle element width, so this should either be enforced with a clear assertion or handled by lowering the descriptor to the usable swizzle atom and looping over the remainder outside the TMA descriptor.
🤖 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/atomic_add.cc` around lines 449 - 459, Update the TMA atomic-add
validation around instruction_dim and desc.smem_box in the visible lowering path
to ensure the global_range extent is divisible by the swizzle element width when
swizzling is enabled. Add a clear assertion for this unsupported remainder case,
or lower the descriptor to complete swizzle atoms and explicitly process any
remainder outside TMA; preserve the existing non-swizzled behavior.
| @pytest.mark.parametrize("dtype", [T.int16, T.int64, T.float64, T.float32x2]) | ||
| def test_tma_atomic_add_rejects_unsupported_dtype(dtype): | ||
| with pytest.raises(Exception, match=rf"TMA atomic add does not support dtype {dtype}.*supported scalar dtypes"): | ||
| with pytest.raises( | ||
| tvm.error.InternalError, | ||
| match=rf"TMA atomic add does not support dtype {dtype}.*supported scalar dtypes", | ||
| ): | ||
| lower_tma_atomic_add(dtype) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check how tilelang dtype classes stringify.
set -euo pipefail
rg -nP -C 8 '__str__|__repr__|class\s+_?DTypeMeta|metaclass=' tilelang/language/dtypes.py
echo "=== other tests interpolating a T.<dtype> into a match pattern ==="
rg -nP -C 3 'match=.*\{dtype\}' testing/python --glob '*.py'Repository: tile-ai/tilelang
Length of output: 154
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== locate dtypes/test files ==="
git ls-files | rg '(^|/)dtypes\.py$|test_tilelang_language_atomic\.py$|testing/python|tilelang' | head -200
echo "=== file existence ==="
for f in tilelang/language/dtypes.py testing/python/language/test_tilelang_language_atomic.py; do
if [ -f "$f" ]; then
echo "FOUND $f"
wc -l "$f"
else
echo "MISSING $.f"
fi
done
echo "=== search __str__/__repr__/metaclass in likely files ==="
rg -n -C 6 '__str__|__repr__|class\s+_?DTypeMeta|metaclass=' $(git ls-files | rg 'dtype|test_tilelang_language_atomic') || trueRepository: tile-ai/tilelang
Length of output: 11365
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== dtypes.py relevant definitions ==="
rg -n -C 12 'class .*DType|def __str__|def __repr__|metaclass|class .*Meta' tilelang/language/dtypes.py
echo "=== test dtype parameterization around TMA test ==="
sed -n '340,410p' testing/python/language/test_tilelang_language_atomic.py
echo "=== imports in test file around top ==="
sed -n '1,80p' testing/python/language/test_tilelang_language_atomic.py
echo "=== precise Python string behavior for a DType instance/class if sources define custom __str__/__repr__ ==="
python3 - <<'PY'
from tilelang.language.dtypes import DType
print("DType is class:", DType)
dt = DType("float32x2")
print("Dtype instance str:", str(dt))
print("Dtype instance repr:", repr(dt))
D = type("DtypeClass", (), {"name": "float32x2"})
print("DtypeClass:", D)
print("DtypeClass str:", str(D))
print("DtypeClass repr:", repr(D))
PYRepository: tile-ai/tilelang
Length of output: 193
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== top dtypes.py ==="
sed -n '1,220p' tilelang/language/dtypes.py
echo "=== any DType definitions/types/strings in tilelang/language/dtypes.py ==="
python3 - <<'PY'
from pathlib import Path
p=Path('tilelang/language/dtypes.py')
text=p.read_text(errors='replace')
for i,line in enumerate(text.splitlines(),1):
if 'DType' in line or 'dtype' in line.lower() or 'str(' in line.lower() or 'repr(' in line.lower() or 'metaclass' in line.lower():
print(f'{i}: {line}')
PY
echo "=== runtime probe if package importable ==="
python3 - <<'PY'
try:
import tilelang.language.dtypes as dpy
print("imports_ok=True")
for name in ['DType','float32x2','int16','int64','float64','float32x2']:
obj=getattr(dpy, name, None)
print(f"{name}: {obj!r} str={str(obj)} repr={repr(obj)}")
dt=dpy.float32x2()
print(f"float32x2 instance: {dt!r} str={str(dt)} repr={repr(dt)}")
except Exception as e:
print(f"imports_ok=False: {type(e).__name__}: {e}")
PY
echo "=== test file relevant sections ==="
sed -n '1,80p' testing/python/language/test_tilelang_language_atomic.py
sed -n '340,410p' testing/python/language/test_tilelang_language_atomic.pyRepository: tile-ai/tilelang
Length of output: 33162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== Python class repr/string behavior ==="
python3 <<'PY'
class dtype:
def __init__(self, value):
self.value = value
T = type("T", (), {})
T.int16 = dtype("int16")
T.float32x2 = dtype("float32x2")
for dtype in [T.int16, T.float32x2]:
print("dtype:", dtype)
print(" str:", str(dtype))
print(" repr:", repr(dtype))
print(" rf-matches:", bool(__import__("re").match(rf"TMA atomic add does not support dtype {dtype}.*supported scalar dtypes", "TMA atomic add does not support dtype int16...")))
PY
echo "=== lower_tma_atomic_add references and TMA error text ==="
rg -n -C 8 'def lower_tma_atomic_add|lower_tma_atomic_add\(|TMA atomic add|does not support dtype|supported scalar dtypes' testing/python/language tilelang --glob '*.py' || trueRepository: tile-ai/tilelang
Length of output: 7909
Interpolate the plain dtype name in the TMA atomic-add rejection test.
dtype values such as T.float32x2 are TileLang dtype instances, not strings. Without an explicit dtype name, rf"... {dtype}" expands to the object representation and cannot match the C++ diagnostic. Use str(dtype) or the dtype value explicitly in the expected message.
🤖 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/language/test_tilelang_language_atomic.py` around lines 383 -
389, Update test_tma_atomic_add_rejects_unsupported_dtype so the regex
interpolates the plain dtype name via str(dtype), matching the C++ diagnostic
for TileLang dtype instances while preserving the existing unsupported-dtype
coverage.
|
@regression-perf |
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
1 participant
Summary
Changes
Validation
./format.shcmake --build build -j$(nproc)python -m pytest -q testing/python/language/test_tilelang_language_atomic.py -k tma_atomic_addpython -m pytest -q testing/python/language/test_tilelang_language_tma_copy.py testing/python/transform/test_tilelang_transform_smem_swizzle_alignment.pyNotes
Summary
uint64, invalid layouts, and 32B swizzle splitting.uint64on unswizzled linear shared storage because the existing 64-bit GEMM K-inner layout cannot be represented by a TensorMap descriptor.C++ style / lint notes
docs/developer_guide/cpp_style.md.Testing