[Refactor] Refactor CUDA atomic helpers - #2001
Merged
Merged
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! 🚀 |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughIntroduces memory-order classification helpers, 16-bit pack/unpack utilities, and centralized PTX inline-assembly atomic-add helpers; refactors fp16/bf16 and vectorized atomic-add implementations to use PTX helpers with scalar fallbacks for older architectures. No public signatures changed. Changes
Sequence Diagram(s)sequenceDiagram
participant Thread
participant PTX_Helper as PTX Helper (inline asm)
participant Scalar_Fallback as Scalar Fallback (per-element atomicAdd)
participant Memory
Thread->>PTX_Helper: if memory_order != relaxed and arch supports PTX
alt PTX path
PTX_Helper->>Memory: perform grouped atomic-add (fp16/bf16/v2/v4)
Memory-->>PTX_Helper: updated value (if returning)
PTX_Helper-->>Thread: result/ack
else scalar fallback
Thread->>Scalar_Fallback: call pair/quad scalar helper
Scalar_Fallback->>Memory: perform element-wise atomicAdd
Memory-->>Scalar_Fallback: per-element results
Scalar_Fallback-->>Thread: combined result
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsTimed out fetching pipeline failures after 30000ms 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 |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/tl_templates/cuda/atomic.h`:
- Around line 56-68: Update the memory-order classifier functions so they match
C++/CUDA semantics and allow the PTX opcode selection to handle consume and
seq_cst specially: change IsReleaseLikeMemoryOrder to exclude
memory_order_consume (only treat memory_order_release as release-like), ensure
IsAcquireMemoryOrder treats memory_order_acquire and memory_order_consume as
acquire-like, and ensure IsAcqRelLikeMemoryOrder does not include
memory_order_seq_cst (only include memory_order_acq_rel); then split handling
for memory_order_consume and memory_order_seq_cst into dedicated branches before
the PTX opcode selection so seq_cst uses the required fence.sc wrapper per the
PTX atomic ABI and consume uses acquire-like semantics.
- Around line 98-116: The BF16 PTX helper AtomicAddPtxBF16 is unguarded and can
emit SM90-only PTX; wrap its definition and any call sites (notably the
return-value caller AtomicAddRet) with the same SM90+ preprocessor guard used
for the void AtomicAdd overload (use `#if` (defined(__CUDA_ARCH_LIST__) &&
(__CUDA_ARCH_LIST__ > 890)) / `#endif`) so PTX instructions like atom.*.bf16 are
only compiled for SM90+, or provide an alternative non-PTX fallback inside the
guarded branch for SM8.x; ensure the symbol AtomicAddPtxBF16 and its use in
AtomicAddRet are enclosed by this guard.
- Around line 118-160: AtomicAddx2 and AtomicAddx2Ret currently call SM90-only
PTX helpers (AtomicAddPtxV2F16 / AtomicAddPtxV2BF16) for non-relaxed orders with
no SM guards, which breaks pre-SM90 (and pre-SM80 for bf16) targets; add
compile-time SM guards and provide fallback ordered implementations: wrap the
SM90 PTX calls in a __CUDA_ARCH__ >= 900 check for v2.f16 (and >= 800/750 as
appropriate for bf16), and in the else branch implement ordered behavior by
performing two lane-wise ordered scalar atomic adds (e.g., call the existing
scalar atomic-add/AtomicAddPtx for each half/bf16 lane or loop with
compare-and-swap to emulate acquire/release semantics) so AtomicAddx2 and
AtomicAddx2Ret work correctly on older SMs. Ensure you modify the functions
AtomicAddx2 / AtomicAddx2Ret to choose the guarded PTX path or the
scalar/emulated fallback, and keep existing relaxed path unchanged.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1bd03053-6bd4-4a91-a4ea-3d0c788f043a
📒 Files selected for processing (1)
src/tl_templates/cuda/atomic.h
Comment on lines
+98
to
+116
| TL_DEVICE void AtomicAddPtxBF16(unsigned short &ret, unsigned long long addr, | ||
| unsigned short val, int memory_order) { | ||
| if (IsReleaseLikeMemoryOrder(memory_order)) { | ||
| asm volatile("atom.release.gpu.global.add.noftz.bf16 %0, [%1], %2;" | ||
| : "=h"(ret) | ||
| : "l"(addr), "h"(val) | ||
| : "memory"); | ||
| } else if (IsAcquireMemoryOrder(memory_order)) { | ||
| asm volatile("atom.acquire.gpu.global.add.noftz.bf16 %0, [%1], %2;" | ||
| : "=h"(ret) | ||
| : "l"(addr), "h"(val) | ||
| : "memory"); | ||
| } else if (IsAcqRelLikeMemoryOrder(memory_order)) { | ||
| asm volatile("atom.acq_rel.gpu.global.add.noftz.bf16 %0, [%1], %2;" | ||
| : "=h"(ret) | ||
| : "l"(addr), "h"(val) | ||
| : "memory"); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's understand the file structure
fd -t f atomic.h | head -20Repository: tile-ai/tilelang
Length of output: 120
🏁 Script executed:
# Get the full content of the atomic.h file to understand the context
wc -l src/tl_templates/cuda/atomic.hRepository: tile-ai/tilelang
Length of output: 94
🏁 Script executed:
# Read the file around the mentioned lines and search for calls to AtomicAddPtxBF16
cat -n src/tl_templates/cuda/atomic.h | head -150Repository: tile-ai/tilelang
Length of output: 6450
🏁 Script executed:
# Search for calls to AtomicAddPtxBF16 and references to SM90
rg -n "AtomicAddPtxBF16|SM90|sm_90" src/tl_templates/cuda/atomic.h -A 2 -B 2Repository: tile-ai/tilelang
Length of output: 825
🏁 Script executed:
# Read context around line 413
cat -n src/tl_templates/cuda/atomic.h | sed -n '400,430p'Repository: tile-ai/tilelang
Length of output: 1443
🏁 Script executed:
# Read context around line 454
cat -n src/tl_templates/cuda/atomic.h | sed -n '440,470p'Repository: tile-ai/tilelang
Length of output: 1535
🏁 Script executed:
# Look for SM90 or CUDA_ARCH guards in the file
rg -n "__CUDA_ARCH|SM90|sm_90|cuda_arch_list" src/tl_templates/cuda/atomic.h -iRepository: tile-ai/tilelang
Length of output: 398
🏁 Script executed:
# Read context around line 388 to see which function has SM90 guard
cat -n src/tl_templates/cuda/atomic.h | sed -n '380,420p'Repository: tile-ai/tilelang
Length of output: 1942
🏁 Script executed:
# Check which functions are wrapped by SM90 guards
cat -n src/tl_templates/cuda/atomic.h | sed -n '375,430p'Repository: tile-ai/tilelang
Length of output: 2584
🏁 Script executed:
# Read more context to see the function with return value
cat -n src/tl_templates/cuda/atomic.h | sed -n '430,480p'Repository: tile-ai/tilelang
Length of output: 2507
🏁 Script executed:
# Check if there's a guard around the return value function that contains line 454
cat -n src/tl_templates/cuda/atomic.h | sed -n '430,468p'Repository: tile-ai/tilelang
Length of output: 1944
🏁 Script executed:
# Check if AtomicAddRet is guarded by any SM version check
cat -n src/tl_templates/cuda/atomic.h | grep -B 5 "TL_DEVICE T1 AtomicAddRet"Repository: tile-ai/tilelang
Length of output: 262
🏁 Script executed:
# Verify the complete structure - check if there are any guards around AtomicAddRet
cat -n src/tl_templates/cuda/atomic.h | sed -n '420,435p'Repository: tile-ai/tilelang
Length of output: 812
🌐 Web query:
PTX atom.add.noftz.bf16 SM90 requirement NVIDIA documentation
💡 Result:
The PTX instruction atom.add.noftz.bf16 (bfloat16 atomic add with no flush-to-zero) requires NVIDIA Hopper architecture (compute capability SM_90) or higher. This is the native hardware atomic add operation for bfloat16 data type in global or shared memory. On pre-SM_90 architectures (e.g., Ampere SM_80), CUDA's atomicAdd for bfloat16 falls back to a software emulation using atomic compare-and-swap (CAS) loops combined with fma.rn.bf16 operations, rather than a direct atomic add instruction. PTX ISA support for bf16 atomics was introduced in version 7.8, coinciding with SM_90.
Citations:
- 1: https://docs.nvidia.com/cuda/archive/11.8.0/parallel-thread-execution/index.html
- 2: https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html
- 3: Feature Request:
tl.atomic_addfor bfloat16 triton-lang/triton#1387 - 4: Feature Request:
tl.atomic_addfor bfloat16 triton-lang/triton#1387 - 5: [BACKEND] Support bf16 global atomic add on Hopper and Ampere triton-lang/triton#2708
- 6: [NVPTX] Select bfloat16 add/mul/sub as fma on SM80 llvm/llvm-project#121065
Guard the scalar BF16 PTX helper to SM90+.
AtomicAddPtxBF16 (lines 98–116) contains PTX instructions like atom.release.gpu.global.add.noftz.bf16 that require SM90+. However, this function is called unconditionally in AtomicAddRet (line 454) when handling non-relaxed memory orders with BF16 types, without any SM version guard. On SM8.x devices, this will attempt to execute unsupported PTX instructions and fail at runtime. The void AtomicAdd overload at line 388 correctly guards this with #if (defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 890)), but the return-value version has no such protection.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/tl_templates/cuda/atomic.h` around lines 98 - 116, The BF16 PTX helper
AtomicAddPtxBF16 is unguarded and can emit SM90-only PTX; wrap its definition
and any call sites (notably the return-value caller AtomicAddRet) with the same
SM90+ preprocessor guard used for the void AtomicAdd overload (use `#if`
(defined(__CUDA_ARCH_LIST__) && (__CUDA_ARCH_LIST__ > 890)) / `#endif`) so PTX
instructions like atom.*.bf16 are only compiled for SM90+, or provide an
alternative non-PTX fallback inside the guarded branch for SM8.x; ensure the
symbol AtomicAddPtxBF16 and its use in AtomicAddRet are enclosed by this guard.
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @SiriusNEO Results
Artifacts
|
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @SiriusNEO Results
Artifacts
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR refactors src/tl_templates/cuda/atomic.h to reduce repeated inline PTX code and make the atomic add paths easier to maintain.
Summary by CodeRabbit