Skip to content

Report unsupported 16-bit integer atomic adds clearly #3370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dajiaohuang wants to merge 7 commits into
base: main
Choose a base branch
from
+46 −5
Open

Report unsupported 16-bit integer atomic adds clearly #3370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b5735ec
Report unsupported 16-bit integer atomic adds
dajiaohuang Oct 1, 2026
d07f2d9
Format unsupported atomic diagnostic
dajiaohuang Oct 1, 2026
af8ec34
test: cover rejected integer atomic return path
dajiaohuang Oct 2, 2026
4f23817
Support bfloat16 scalar inputs in packed and Cython calls
dajiaohuang Oct 3, 2026
174d67b
Load packed bfloat16 scalar values through the FFI double field
dajiaohuang Oct 3, 2026
1ad203c
Merge scalar bfloat16 correction for native validation
dajiaohuang Oct 3, 2026
d9f5dae
fix: keep atomic diagnostic PR scoped to int16
dajiaohuang Oct 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 22 additions & 5 deletions src/tl_templates/cuda/atomic.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,12 @@ TL_DEVICE void AtomicAddx4Scalar(T *ref, T x, T y, T z, T w) {
atomicAdd(ref + 3, w);
}

template <typename T> TL_DEVICE void CheckAtomicAddDestinationType() {
static_assert(
!std::is_same_v<T, short> && !std::is_same_v<T, unsigned short>,
"CUDA atomic_add does not support int16 or uint16 destinations");
}

TL_DEVICE float2 AtomicAddx2ScalarRet(float *ref, float2 add_val) {
float2 ret;
ret.x = atomicAdd(ref + 0, add_val.x);
Expand Down Expand Up @@ -472,8 +478,11 @@ template <typename T1, typename T2>
TL_DEVICE void AtomicAdd(T1 *address, T2 val,
int memory_order = int(cuda::memory_order_relaxed)) {
using NT1 = typename normalize_atomic_type<T1>::type;
if constexpr (std::is_same_v<NT1, half> ||
std::is_same_v<NT1, __nv_bfloat16>) {
if constexpr (std::is_same_v<NT1, short> ||
std::is_same_v<NT1, unsigned short>) {
tl_atomic_detail::CheckAtomicAddDestinationType<NT1>();
} else if constexpr (std::is_same_v<NT1, half> ||
std::is_same_v<NT1, __nv_bfloat16>) {
if (tl_atomic_detail::IsRelaxedMemoryOrder(memory_order)) {
atomicAdd(reinterpret_cast<NT1 *>(address), static_cast<NT1>(val));
} else {
Expand Down Expand Up @@ -508,16 +517,24 @@ template <typename T1, typename T2>
TL_DEVICE void AtomicAdd(T1 *address, T2 val,
int memory_order = int(cuda::memory_order_relaxed)) {
using NT1 = typename normalize_atomic_type<T1>::type;
(void)memory_order;
atomicAdd(reinterpret_cast<NT1 *>(address), cuda_cast<NT1>(val));
if constexpr (std::is_same_v<NT1, short> ||
std::is_same_v<NT1, unsigned short>) {
tl_atomic_detail::CheckAtomicAddDestinationType<NT1>();
} else {
(void)memory_order;
atomicAdd(reinterpret_cast<NT1 *>(address), cuda_cast<NT1>(val));
}
}
#endif

template <typename T1, typename T2>
TL_DEVICE T1 AtomicAddRet(T1 *address, T2 val,
int memory_order = int(cuda::memory_order_relaxed)) {
using NT1 = typename normalize_atomic_type<T1>::type;
if constexpr (std::is_same_v<NT1, bfloat16_t>) {
if constexpr (std::is_same_v<NT1, short> ||
std::is_same_v<NT1, unsigned short>) {
tl_atomic_detail::CheckAtomicAddDestinationType<NT1>();
} else if constexpr (std::is_same_v<NT1, bfloat16_t>) {
// Pre-SM80 only: cuda::atomic_ref has no fetch_add for bfloat16_t, so use
// the atomicAdd overload above. Memory order is dropped, as in AtomicAdd.
(void)memory_order;
Expand Down
24 changes: 24 additions & 0 deletions testing/python/language/test_tilelang_language_atomic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,6 +665,30 @@ def test_atomic_add():
run_atomic_add(8, 128, 128, 32, 32)


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("dtype", [T.int16, T.uint16])
def test_atomic_add_rejects_16bit_integer_destination(dtype):
@T.prim_func
def atomic_add_16bit_integer(src: T.Tensor((1,), dtype), dst: T.Tensor((1,), dtype)):
with T.Kernel(1, threads=1):
T.atomic_add(dst[0], src[0])

with pytest.raises(RuntimeError, match="CUDA atomic_add does not support int16 or uint16 destinations"):
tilelang.compile(atomic_add_16bit_integer,>


@tilelang.testing.requires_cuda
@pytest.mark.parametrize("dtype", [T.int16, T.uint16])
def test_atomic_add_return_prev_rejects_16bit_integer_destination(dtype):
@T.prim_func
def atomic_add_16bit_integer(src: T.Tensor((1,), dtype), dst: T.Tensor((1,), dtype), prev: T.Tensor((1,), dtype)):
with T.Kernel(1, threads=1):
prev[0] = T.atomic_add(dst[0], src[0], return_prev=True)

with pytest.raises(RuntimeError, match="CUDA atomic_add does not support int16 or uint16 destinations"):
tilelang.compile(atomic_add_16bit_integer,>


@tilelang.testing.requires_cuda
@tilelang.testing.requires_cuda_compute_version_ge(9, 0)
def test_atomic_add_auto_vectorized():
Expand Down
Loading