Skip to content

[Refactor] Move AtomicAdd Vectorization to VectorizeLoop Pass #1677

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

Merged
LeiWang1999 merged 6 commits into from
Jan 16, 2026
Merged

[Refactor] Move AtomicAdd Vectorization to VectorizeLoop Pass #1677

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
639e8e4
move to vectorize_loop
LeiWang1999 Jan 15, 2026
dcff47e
fix
LeiWang1999 Jan 16, 2026
53f4a5d
lint fix
LeiWang1999 Jan 16, 2026
94370c3
Refactor atomic add functions to improve type handling and reduce red…
LeiWang1999 Jan 16, 2026
f1bf450
add related test
LeiWang1999 Jan 16, 2026
173fe4d
Remove atomic add vectorization implementation
LeiWang1999 Jan 16, 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
194 changes: 181 additions & 13 deletions src/tl_templates/cuda/atomic.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -327,8 +327,8 @@ TL_DEVICE T1 AtomicAddRet(T1 *address, T2 val,
}
}

// TODO add memory_order for vectorized atomic add
TL_DEVICE void AtomicAddx2(half_t *ref, half_t *val,
template <typename src_type>
TL_DEVICE void AtomicAddx2(half_t *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(reinterpret_cast<half2 *>(ref),
Expand Down Expand Up @@ -374,8 +374,9 @@ TL_DEVICE void AtomicAddx2(half_t *ref, half_t *val,
}
}

template <typename src_type>
TL_DEVICE half2
AtomicAddx2Ret(half_t *ref, half_t *val,
AtomicAddx2Ret(half_t *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(reinterpret_cast<half2 *>(ref),
Expand Down Expand Up @@ -419,7 +420,8 @@ AtomicAddx2Ret(half_t *ref, half_t *val,
}

#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 750))
TL_DEVICE void AtomicAddx2(bfloat16_t *ref, bfloat16_t *val,
template <typename src_type>
TL_DEVICE void AtomicAddx2(bfloat16_t *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(
Expand Down Expand Up @@ -458,8 +460,9 @@ TL_DEVICE void AtomicAddx2(bfloat16_t *ref, bfloat16_t *val,
}
}

template <typename src_type>
TL_DEVICE __nv_bfloat162
AtomicAddx2Ret(bfloat16_t *ref, bfloat16_t *val,
AtomicAddx2Ret(bfloat16_t *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(
Expand Down Expand Up @@ -502,7 +505,8 @@ AtomicAddx2Ret(bfloat16_t *ref, bfloat16_t *val,
#endif

#if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ >= 900))
TL_DEVICE void AtomicAddx2(float *ref, float *val,
template <typename src_type>
TL_DEVICE void AtomicAddx2(float *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(reinterpret_cast<float2 *>(ref),
Expand Down Expand Up @@ -532,8 +536,37 @@ TL_DEVICE void AtomicAddx2(float *ref, float *val,
}
}

TL_DEVICE void AtomicAddx2(float *ref, float2 val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(reinterpret_cast<float2 *>(ref), val);
} else {
unsigned long long ref_addr = reinterpret_cast<unsigned long long>(ref);
float2 ret_val;
if (memory_order == int(cuda::memory_order_release) ||
memory_order == int(cuda::memory_order_consume)) {
asm volatile("atom.release.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
} else if (memory_order == int(cuda::memory_order_acquire)) {
asm volatile("atom.acquire.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
} else if (memory_order == int(cuda::memory_order_acq_rel) ||
memory_order == int(cuda::memory_order_seq_cst)) {
asm volatile("atom.acq_rel.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
}
}
}

template <typename src_type>
TL_DEVICE float2
AtomicAddx2Ret(float *ref, float *val,
AtomicAddx2Ret(float *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(reinterpret_cast<float2 *>(ref),
Expand Down Expand Up @@ -564,7 +597,39 @@ AtomicAddx2Ret(float *ref, float *val,
}
}

TL_DEVICE void AtomicAddx4(float *ref, float *val,
template <typename src_type>
TL_DEVICE float2
AtomicAddx2Ret(float *ref, float2 val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(reinterpret_cast<float2 *>(ref), static_cast<float2>(val));
} else {
unsigned long long ref_addr = reinterpret_cast<unsigned long long>(ref);
float2 ret_val;
if (memory_order == int(cuda::memory_order_release) ||
memory_order == int(cuda::memory_order_consume)) {
asm volatile("atom.release.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
} else if (memory_order == int(cuda::memory_order_acquire)) {
asm volatile("atom.acquire.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
} else if (memory_order == int(cuda::memory_order_acq_rel) ||
memory_order == int(cuda::memory_order_seq_cst)) {
asm volatile("atom.acq_rel.gpu.global.add.v2.f32 {%0,%1}, [%2], {%3,%4};"
: "=f"(ret_val.x), "=f"(ret_val.y)
: "l"(ref_addr), "f"(val.x), "f"(val.y)
: "memory");
}
return ret_val;
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

template <typename src_dtype, typename dst_dtype>
TL_DEVICE void AtomicAddx4(dst_dtype *ref, src_dtype *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(reinterpret_cast<float4 *>(ref),
Expand Down Expand Up @@ -606,8 +671,47 @@ TL_DEVICE void AtomicAddx4(float *ref, float *val,
}
}

template <typename dst_dtype>
TL_DEVICE void AtomicAddx4(dst_dtype *ref, float4 val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
atomicAdd(reinterpret_cast<float4 *>(ref), val);
} else {
// Since atomicAdd does not support memory order, atomic_ref does not
// support vectorized atomic operation we can only inline ptx code here
// Note: Vectorized atomic operations only support global space
unsigned long long ref_addr = reinterpret_cast<unsigned long long>(ref);
float4 ret_val;
if (memory_order == int(cuda::memory_order_release) ||
memory_order == int(cuda::memory_order_consume)) {
asm volatile(
"atom.release.gpu.global.add.v4.f32 {%0,%1,%2,%3}, [%4], "
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
} else if (memory_order == int(cuda::memory_order_acquire)) {
asm volatile(
"atom.acquire.gpu.global.add.v4.f32 {%0,%1,%2,%3}, [%4], "
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
} else if (memory_order == int(cuda::memory_order_acq_rel) ||
memory_order == int(cuda::memory_order_seq_cst)) {
asm volatile(
"atom.acq_rel.gpu.global.add.v4.f32 {%0,%1,%2,%3}, [%4], "
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
}
}
}

template <typename src_dtype, typename dst_dtype>
TL_DEVICE float4
AtomicAddx4Ret(float *ref, float *val,
AtomicAddx4Ret(dst_dtype *ref, src_dtype *val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(reinterpret_cast<float4 *>(ref),
Expand Down Expand Up @@ -646,17 +750,56 @@ AtomicAddx4Ret(float *ref, float *val,
return ret_val;
}
}

template <typename dst_dtype>
TL_DEVICE float4
AtomicAddx4Ret(dst_dtype *ref, float4 val,
int memory_order = int(cuda::memory_order_relaxed)) {
if (memory_order == int(cuda::memory_order_relaxed)) {
return atomicAdd(reinterpret_cast<float4 *>(ref), val);
} else {
unsigned long long ref_addr = reinterpret_cast<unsigned long long>(ref);
float4 ret_val;
if (memory_order == int(cuda::memory_order_release) ||
memory_order == int(cuda::memory_order_consume)) {
asm volatile(
"atom.global.gpu.release.add.v4.f32 {%0,%1,%2,%3}, [%4], "

Copilot AI Jan 15, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PTX instruction format is incorrect. The order should be "atom.release.gpu.global.add.v4.f32" not "atom.global.gpu.release.add.v4.f32". The memory order (release) should come before the scope (global), matching the format used in lines 624, 631, and 639 for the similar overload.

Copilot uses AI. Check for mistakes.
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
} else if (memory_order == int(cuda::memory_order_acquire)) {
asm volatile(
"atom.global.gpu.acquire.add.v4.f32 {%0,%1,%2,%3}, [%4], "

Copilot AI Jan 15, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PTX instruction format is incorrect. The order should be "atom.acquire.gpu.global.add.v4.f32" not "atom.global.gpu.acquire.add.v4.f32". The memory order (acquire) should come before the scope (global), matching the format used in line 630 for the similar overload.

Copilot uses AI. Check for mistakes.
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
} else if (memory_order == int(cuda::memory_order_acq_rel) ||
memory_order == int(cuda::memory_order_seq_cst)) {
asm volatile(
"atom.global.gpu.acq_rel.add.v4.f32 {%0,%1,%2,%3}, [%4], "

Copilot AI Jan 15, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PTX instruction format is incorrect. The order should be "atom.acq_rel.gpu.global.add.v4.f32" not "atom.global.gpu.acq_rel.add.v4.f32". The memory order (acq_rel) should come before the scope (global), matching the format used in lines 638-639 for the similar overload.

Suggested change
"atom.global.gpu.acq_rel.add.v4.f32 {%0,%1,%2,%3}, [%4], "
"atom.acq_rel.gpu.global.add.v4.f32 {%0,%1,%2,%3}, [%4], "

Copilot uses AI. Check for mistakes.
"{%5,%6,%7,%8};"
: "=f"(ret_val.x), "=f"(ret_val.y), "=f"(ret_val.z), "=f"(ret_val.w)
: "l"(ref_addr), "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w)
: "memory");
}
return ret_val;
}
}
#else
TL_DEVICE void AtomicAddx2(float *ref, float *val,
template <typename src_type>
TL_DEVICE void AtomicAddx2(float *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
float2 add_val = *reinterpret_cast<float2 *>(val);
atomicAdd(ref + 0, add_val.x);
atomicAdd(ref + 1, add_val.y);
}

template <typename src_type>
TL_DEVICE float2
AtomicAddx2Ret(float *ref, float *val,
AtomicAddx2Ret(float *ref, src_type *val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
float2 add_val = *reinterpret_cast<float2 *>(val);
Expand All @@ -666,7 +809,8 @@ AtomicAddx2Ret(float *ref, float *val,
return ret;
}

TL_DEVICE void AtomicAddx4(float *ref, float *val,
template <typename src_dtype, typename dst_dtype>
TL_DEVICE void AtomicAddx4(dst_dtype *ref, src_dtype *val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
float4 add_val = *reinterpret_cast<float4 *>(val);
Expand All @@ -676,8 +820,19 @@ TL_DEVICE void AtomicAddx4(float *ref, float *val,
atomicAdd(ref + 3, add_val.w);
}

template <typename dst_dtype>
TL_DEVICE void AtomicAddx4(dst_dtype *ref, float4 val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
atomicAdd(ref + 0, val.x);
atomicAdd(ref + 1, val.y);
atomicAdd(ref + 2, val.z);
atomicAdd(ref + 3, val.w);
}

template <typename src_dtype, typename dst_dtype>
TL_DEVICE float4
AtomicAddx4Ret(float *ref, float *val,
AtomicAddx4Ret(dst_dtype *ref, src_dtype *val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
float4 add_val = *reinterpret_cast<float4 *>(val);
Expand All @@ -688,6 +843,19 @@ AtomicAddx4Ret(float *ref, float *val,
ret.w = atomicAdd(ref + 3, add_val.w);
return ret;
}

template <typename dst_dtype>
TL_DEVICE float4
AtomicAddx4Ret(dst_dtype *ref, float4 val,
int memory_order = int(cuda::memory_order_relaxed)) {
(void)memory_order;
float4 ret;
ret.x = atomicAdd(ref + 0, val.x);
ret.y = atomicAdd(ref + 1, val.y);
ret.z = atomicAdd(ref + 2, val.z);
ret.w = atomicAdd(ref + 3, val.w);
return ret;
}
#endif

template <typename T> TL_DEVICE T AtomicLoad(T *ref, int memory_order) {
Expand Down
6 changes: 1 addition & 5 deletions src/transform/loop_partition.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,11 +297,7 @@ Stmt LowerParallelLoop(For loop, const Fragment &loop_layout, Var thread_var,
if (should_vectorize) {
result_loop = VectorizeLoop(result_loop, saved_analyzer.get(), layout_map);
}

// Step 3: Vectorize atomic add operations
result_loop = VectorizeAtomicAdd(result_loop);

// Step 4: Wrap with predicate if provided and this is a parallel loop
// Step 3: Wrap with predicate if provided and this is a parallel loop
if (predicate.defined() && parallel_loop) {
return IfThenElse(predicate.value(), result_loop);
}
Expand Down
Loading
Loading