Repository navigation
[NVPTX] Respect FTZ flag when lowering atomicrmw fadd. - #200732
Conversation
|
@llvm/pr-subscribers-backend-nvptx Author: Justin Lebar (jlebar) ChangesPreviously we unconditionally lowered LLVM atomicrmw fadd to PTX Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Full diff: https://github.com/llvm/llvm-project/pull/200732.diff 2 Files Affected:
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index d1d01089aa49e..f45491f66cca8 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -7467,22 +7467,53 @@ NVPTXTargetLowering::AtomicExpansionKind
NVPTXTargetLowering::shouldExpandAtomicRMWInIR(const AtomicRMWInst *AI) const {
Type *Ty = AI->getValOperand()->getType();
- if (AI->isFloatingPointOperation()) {
- if (AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
- if (Ty->isHalfTy() && STI.getSmVersion() >= 70 &&
- STI.getPTXVersion() >= 63)
- return AtomicExpansionKind::None;
- if (Ty->isBFloatTy() && STI.getSmVersion() >= 90 &&
- STI.getPTXVersion() >= 78)
- return AtomicExpansionKind::None;
- if (Ty->isFloatTy())
- return AtomicExpansionKind::None;
- if (Ty->isDoubleTy() && STI.hasAtomAddF64())
- return AtomicExpansionKind::None;
+ // Try to lower LLVM atomicrmw fadd to PTX atomic.add. This is complicated
+ // by the weird FTZ behavior PTX atom.add has:
+ // - atom.add.f32 on global memory flushes denormals
+ // - atom.add.f32 on shared memory does not flush denormals
+ // - atom.add.f16 and atomic.add.bf16 never flush denormals
+ //
+ // We lower to atom.add only if the function's FTZ behavior matches that of
+ // atom.add; otherwise, we lower to a CAS loop. But we always allow
+ // atomic.add.bf16; even though it never flushes denormals, we never flush
+ // bf16 denormals when doing regular arithmetic, even when FTZ is enabled.
+ if (AI->isFloatingPointOperation() &&
+ AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
+ const bool FTZ =
+ AI->getFunction()->getDenormalMode(APFloat::IEEEsingle()).Output ==
+ DenormalMode::PreserveSign;
+
+ if (Ty->isFloatTy()) {
+ switch (AI->getPointerAddressSpace()) {
+ case llvm::ADDRESS_SPACE_GLOBAL:
+ if (FTZ)
+ return AtomicExpansionKind::None;
+ break;
+ case llvm::ADDRESS_SPACE_SHARED:
+ case llvm::ADDRESS_SPACE_SHARED_CLUSTER:
+ if (!FTZ)
+ return AtomicExpansionKind::None;
+ break;
+ }
}
- return AtomicExpansionKind::CmpXChg;
+
+ if (Ty->isHalfTy() && !FTZ && STI.getSmVersion() >= 70 &&
+ STI.getPTXVersion() >= 63)
+ return AtomicExpansionKind::None;
+
+ if (Ty->isBFloatTy() && STI.getSmVersion() >= 90 &&
+ STI.getPTXVersion() >= 78)
+ return AtomicExpansionKind::None;
+
+ if (Ty->isDoubleTy() && STI.hasAtomAddF64())
+ return AtomicExpansionKind::None;
}
+ // All other atomic fp ops expand to a CAS loop.
+ // TODO: Implement other atomic ops!
+ if (AI->isFloatingPointOperation())
+ return AtomicExpansionKind::CmpXChg;
+
assert(Ty->isIntegerTy() && "Ty should be integer at this point");
const unsigned BitWidth = cast<IntegerType>(Ty)->getBitWidth();
diff --git a/llvm/test/CodeGen/NVPTX/atomicrmw-sm90.ll b/llvm/test/CodeGen/NVPTX/atomicrmw-sm90.ll
index 52dbb2831e27c..bbb05920bcb4e 100644
--- a/llvm/test/CodeGen/NVPTX/atomicrmw-sm90.ll
+++ b/llvm/test/CodeGen/NVPTX/atomicrmw-sm90.ll
@@ -1739,14 +1739,25 @@ define i64 @usub_sat_acq_rel_i64_global_cta(ptr addrspace(1) %addr, i64 %val) {
define float @fadd_acq_rel_float_global_cta(ptr addrspace(1) %addr, float %val) {
; SM90-LABEL: fadd_acq_rel_float_global_cta(
; SM90: {
-; SM90-NEXT: .reg .b32 %r<3>;
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b32 %r<5>;
; SM90-NEXT: .reg .b64 %rd<2>;
; SM90-EMPTY:
; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b32 %r2, [fadd_acq_rel_float_global_cta_param_1];
; SM90-NEXT: ld.param.b64 %rd1, [fadd_acq_rel_float_global_cta_param_0];
-; SM90-NEXT: ld.param.b32 %r1, [fadd_acq_rel_float_global_cta_param_1];
-; SM90-NEXT: atom.acq_rel.cta.global.add.f32 %r2, [%rd1], %r1;
-; SM90-NEXT: st.param.b32 [func_retval0], %r2;
+; SM90-NEXT: fence.release.cta;
+; SM90-NEXT: ld.relaxed.cta.global.b32 %r4, [%rd1];
+; SM90-NEXT: $L__BB60_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: add.rn.f32 %r3, %r4, %r2;
+; SM90-NEXT: atom.relaxed.cta.global.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NEXT: setp.ne.b32 %p1, %r1, %r4;
+; SM90-NEXT: mov.b32 %r4, %r1;
+; SM90-NEXT: @%p1 bra $L__BB60_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: fence.acquire.cta;
+; SM90-NEXT: st.param.b32 [func_retval0], %r1;
; SM90-NEXT: ret;
%retval = atomicrmw fadd ptr addrspace(1) %addr, float %val syncscope("block") acq_rel
ret float %retval
@@ -2981,3 +2992,231 @@ define i32 @nand_seq_cst_i32_global_cta(ptr addrspace(1) %addr, i32 %val) {
ret i32 %retval
}
+
+; Regression test for bug 033. atom.add.f32's subnormal flushing is
+; address-space dependent: global memory flushes (FTZ), shared memory
+; preserves subnormals (IEEE), and generic memory is indeterminate. We may
+; lower atomicrmw fadd float to native atom.add.f32 only when its flush
+; behavior provably matches the function's f32 denormal output mode;
+; otherwise we expand to a CAS loop using the non-flushing add.rn.f32.
+
+; global + FTZ (preserve-sign): native atom.add.f32 flushes, which matches.
+define float @fadd_f32_global_ftz_native(ptr addrspace(1) %p, float %v) #0 {
+; SM90-LABEL: fadd_f32_global_ftz_native(
+; SM90: {
+; SM90-NEXT: .reg .b32 %r<3>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_global_ftz_native_param_0];
+; SM90-NEXT: ld.param.b32 %r1, [fadd_f32_global_ftz_native_param_1];
+; SM90-NEXT: atom.relaxed.sys.global.add.f32 %r2, [%rd1], %r1;
+; SM90-NEXT: st.param.b32 [func_retval0], %r2;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(1) %p, float %v monotonic
+ ret float %r
+}
+
+; global + ieee: native flushes but the function wants IEEE, so expand.
+define float @fadd_f32_global_ieee_expand(ptr addrspace(1) %p, float %v) {
+; SM90-LABEL: fadd_f32_global_ieee_expand(
+; SM90: {
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b32 %r<5>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b32 %r2, [fadd_f32_global_ieee_expand_param_1];
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_global_ieee_expand_param_0];
+; SM90-NEXT: ld.relaxed.sys.global.b32 %r4, [%rd1];
+; SM90-NEXT: $L__BB101_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: add.rn.f32 %r3, %r4, %r2;
+; SM90-NEXT: atom.relaxed.sys.global.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NEXT: setp.ne.b32 %p1, %r1, %r4;
+; SM90-NEXT: mov.b32 %r4, %r1;
+; SM90-NEXT: @%p1 bra $L__BB101_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: st.param.b32 [func_retval0], %r1;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(1) %p, float %v monotonic
+ ret float %r
+}
+
+; shared + ieee: native atom.add.f32 preserves subnormals, which matches.
+define float @fadd_f32_shared_ieee_native(ptr addrspace(3) %p, float %v) {
+; SM90-LABEL: fadd_f32_shared_ieee_native(
+; SM90: {
+; SM90-NEXT: .reg .b32 %r<3>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_shared_ieee_native_param_0];
+; SM90-NEXT: ld.param.b32 %r1, [fadd_f32_shared_ieee_native_param_1];
+; SM90-NEXT: atom.relaxed.sys.shared.add.f32 %r2, [%rd1], %r1;
+; SM90-NEXT: st.param.b32 [func_retval0], %r2;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(3) %p, float %v monotonic
+ ret float %r
+}
+
+; shared + FTZ (preserve-sign): native preserves but the function wants
+; flushing, so expand.
+define float @fadd_f32_shared_ftz_expand(ptr addrspace(3) %p, float %v) #0 {
+; SM90-LABEL: fadd_f32_shared_ftz_expand(
+; SM90: {
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b32 %r<5>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b32 %r2, [fadd_f32_shared_ftz_expand_param_1];
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_shared_ftz_expand_param_0];
+; SM90-NEXT: ld.relaxed.sys.shared.b32 %r4, [%rd1];
+; SM90-NEXT: $L__BB103_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: add.rn.ftz.f32 %r3, %r4, %r2;
+; SM90-NEXT: atom.relaxed.sys.shared.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NEXT: setp.ne.b32 %p1, %r1, %r4;
+; SM90-NEXT: mov.b32 %r4, %r1;
+; SM90-NEXT: @%p1 bra $L__BB103_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: st.param.b32 [func_retval0], %r1;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(3) %p, float %v monotonic
+ ret float %r
+}
+
+; generic + ieee: flush behavior is indeterminate, so expand.
+define float @fadd_f32_generic_ieee_expand(ptr %p, float %v) {
+; SM90-LABEL: fadd_f32_generic_ieee_expand(
+; SM90: {
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b32 %r<5>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b32 %r2, [fadd_f32_generic_ieee_expand_param_1];
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_generic_ieee_expand_param_0];
+; SM90-NEXT: ld.relaxed.sys.b32 %r4, [%rd1];
+; SM90-NEXT: $L__BB104_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: add.rn.f32 %r3, %r4, %r2;
+; SM90-NEXT: atom.relaxed.sys.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NEXT: setp.ne.b32 %p1, %r1, %r4;
+; SM90-NEXT: mov.b32 %r4, %r1;
+; SM90-NEXT: @%p1 bra $L__BB104_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: st.param.b32 [func_retval0], %r1;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr %p, float %v monotonic
+ ret float %r
+}
+
+; generic + FTZ (preserve-sign): flush behavior is indeterminate, so expand.
+define float @fadd_f32_generic_ftz_expand(ptr %p, float %v) #0 {
+; SM90-LABEL: fadd_f32_generic_ftz_expand(
+; SM90: {
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b32 %r<5>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b32 %r2, [fadd_f32_generic_ftz_expand_param_1];
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f32_generic_ftz_expand_param_0];
+; SM90-NEXT: ld.relaxed.sys.b32 %r4, [%rd1];
+; SM90-NEXT: $L__BB105_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: add.rn.ftz.f32 %r3, %r4, %r2;
+; SM90-NEXT: atom.relaxed.sys.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NEXT: setp.ne.b32 %p1, %r1, %r4;
+; SM90-NEXT: mov.b32 %r4, %r1;
+; SM90-NEXT: @%p1 bra $L__BB105_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: st.param.b32 [func_retval0], %r1;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr %p, float %v monotonic
+ ret float %r
+}
+
+; atom.add.f16 requires .noftz (always preserves subnormals). f16 flushing
+; follows the f32 denormal mode (useF32FTZ), so under FTZ the native
+; instruction would diverge from the function's other f16 ops; expand instead.
+
+; f16 + FTZ: expand to a CAS loop whose add.rn.ftz.f16 flushes.
+define half @fadd_f16_ftz_expand(ptr addrspace(1) %p, half %v) #0 {
+; SM90-LABEL: fadd_f16_ftz_expand(
+; SM90: {
+; SM90-NEXT: .reg .pred %p<2>;
+; SM90-NEXT: .reg .b16 %rs<4>;
+; SM90-NEXT: .reg .b32 %r<15>;
+; SM90-NEXT: .reg .b64 %rd<3>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b16 %rs1, [fadd_f16_ftz_expand_param_1];
+; SM90-NEXT: ld.param.b64 %rd2, [fadd_f16_ftz_expand_param_0];
+; SM90-NEXT: and.b64 %rd1, %rd2, -4;
+; SM90-NEXT: cvt.u32.u64 %r4, %rd2;
+; SM90-NEXT: and.b32 %r5, %r4, 3;
+; SM90-NEXT: shl.b32 %r1, %r5, 3;
+; SM90-NEXT: mov.b32 %r6, 65535;
+; SM90-NEXT: shl.b32 %r7, %r6, %r1;
+; SM90-NEXT: not.b32 %r2, %r7;
+; SM90-NEXT: ld.relaxed.sys.global.b32 %r14, [%rd1];
+; SM90-NEXT: $L__BB106_1: // %atomicrmw.start
+; SM90-NEXT: // =>This Inner Loop Header: Depth=1
+; SM90-NEXT: shr.u32 %r8, %r14, %r1;
+; SM90-NEXT: cvt.u16.u32 %rs2, %r8;
+; SM90-NEXT: add.rn.ftz.f16 %rs3, %rs2, %rs1;
+; SM90-NEXT: cvt.u32.u16 %r9, %rs3;
+; SM90-NEXT: shl.b32 %r10, %r9, %r1;
+; SM90-NEXT: and.b32 %r11, %r14, %r2;
+; SM90-NEXT: or.b32 %r12, %r11, %r10;
+; SM90-NEXT: atom.relaxed.sys.global.cas.b32 %r3, [%rd1], %r14, %r12;
+; SM90-NEXT: setp.ne.b32 %p1, %r3, %r14;
+; SM90-NEXT: mov.b32 %r14, %r3;
+; SM90-NEXT: @%p1 bra $L__BB106_1;
+; SM90-NEXT: // %bb.2: // %atomicrmw.end
+; SM90-NEXT: shr.u32 %r13, %r3, %r1;
+; SM90-NEXT: st.param.b16 [func_retval0], %r13;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(1) %p, half %v monotonic
+ ret half %r
+}
+
+; f16 + ieee: native atom.add.noftz.f16 preserves subnormals, which matches.
+define half @fadd_f16_ieee_native(ptr addrspace(1) %p, half %v) {
+; SM90-LABEL: fadd_f16_ieee_native(
+; SM90: {
+; SM90-NEXT: .reg .b16 %rs<3>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_f16_ieee_native_param_0];
+; SM90-NEXT: ld.param.b16 %rs1, [fadd_f16_ieee_native_param_1];
+; SM90-NEXT: atom.relaxed.sys.global.add.noftz.f16 %rs2, [%rd1], %rs1;
+; SM90-NEXT: st.param.b16 [func_retval0], %rs2;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(1) %p, half %v monotonic
+ ret half %r
+}
+
+; bf16 + FTZ: NVPTX has no flushing bf16 add, so native atom.add.noftz.bf16
+; always matches and stays selected even in FTZ mode.
+define bfloat @fadd_bf16_ftz_native(ptr addrspace(1) %p, bfloat %v) #0 {
+; SM90-LABEL: fadd_bf16_ftz_native(
+; SM90: {
+; SM90-NEXT: .reg .b16 %rs<3>;
+; SM90-NEXT: .reg .b64 %rd<2>;
+; SM90-EMPTY:
+; SM90-NEXT: // %bb.0:
+; SM90-NEXT: ld.param.b64 %rd1, [fadd_bf16_ftz_native_param_0];
+; SM90-NEXT: ld.param.b16 %rs1, [fadd_bf16_ftz_native_param_1];
+; SM90-NEXT: atom.relaxed.sys.global.add.noftz.bf16 %rs2, [%rd1], %rs1;
+; SM90-NEXT: st.param.b16 [func_retval0], %rs2;
+; SM90-NEXT: ret;
+ %r = atomicrmw fadd ptr addrspace(1) %p, bfloat %v monotonic
+ ret bfloat %r
+}
+
+attributes #0 = { denormal_fpenv(float: preservesign) }
|
fea9097 to
5b737bc
Compare
|
I suppose the semantics are what they are and we must honor them but I do worry a little bit about performance. In my experience often when people enable FTZ they just want their code to be faster and they aren't to particular about needing ever operation to do FTZ. With this patch we're going to start generating lots of loops where previously we had a single atomic instruction in order to handle a case we don't have any evidence any real user code cares about. @jlebar what do you think? |
|
Yeah, this could regress the performance of some programs. IME code generators like Triton often use inline asm when they Care Deeply about what exactly gets generated. Today that's kind of the only option you have if you want to express that a particular operation gets lowered to a particular set of PTX instructions. If the net effect of this patch is that we push people towards using inline asm for their atom.add's, I don't think that's the wrong direction (at least, until LLVM grows the ability to specify these operations more surgically). Also, this one was particularly bad because even when FTZ was disabled we'd generate atom.add.f32, silently introducing FTZ behavior. If we wanted "moar FTZ" we could enable atom.add on the generic and shared address spaces when function-level FTZ is enabled, but it's a semantics violation, and I am not convinced that's the right thing to do.
We also don't have evidence that real user code is slowed down by these additional loops, though. I'm reminded of an XLA:GPU change we made while I was at Waymo that upgraded the precision of some mathematical operations. Folks argued that we should do this because it was faster and a strict improvement to the mathematical accuracy. My argument was that we shouldn't do it because it was semantically incorrect. We ended up guarding it with a flag, because we didn't have evidence from real users either way. Some months later we found real programs that were broken by the upgraded precision. |
|
@akshayrdeodhar for vis. |
| // | ||
| // We lower to atom.add only if the function's FTZ behavior matches that of | ||
| // atom.add; otherwise, we lower to a CAS loop. But we always allow | ||
| // atomic.add.bf16; even though it never flushes denormals, we never flush |
There was a problem hiding this comment.
@akshayrdeodhar atom.add.bf16 is a 16-bit wide atom, so i'd expect that to lower to a cas-loop (only the two element vector version or wider would lower to atom.add.bf16x2).
There was a problem hiding this comment.
@akshayrdeodhar atom.add.bf16 is a 16-bit wide atom, so i'd expect that to lower to a cas-loop (only the two element vector version or wider would lower to atom.add.bf16x2).
I don't understand wym, can you rephrase? PTX has atom.add.bf16; are you saying you don't want to use it?
atom.add.noftz.bf16 and atom.add.noftz.bf16x2 introduced in PTX ISA 7.8.
There was a problem hiding this comment.
(Or are you saying that atom.add.bf16 will lower to a CAS loop inside the SASS, because it's a 16-bit atom? I believe you. As a result do you not want to use the PTX atomic?)
There was a problem hiding this comment.
I agree, I looked at the SASS dump for atom.add.bf16, it does generate a cas-loop.
There was a problem hiding this comment.
I agree, I looked at the SASS dump for atom.add.bf16, it does generate a cas-loop.
So is the conclusion that we should not generate atom.add.bf16, or any other PTX atom.add's that generate a cas loop? Instead have clang generate the CAS loop itself?
There was a problem hiding this comment.
I propose that if you want to stop emitting atomic.add.f16 and .bf16 (because these are not native ops), we do that in a separate PR.
There was a problem hiding this comment.
@akshayrdeodhar atom.add.bf16 should already not have been emitted, because we set the smallest atomic width supported to 32-bit, so that AtomicExpand always uses a cas-loop for anything smaller.
Was NVPTX accidentally opting out of that here for bf16 and float?
Or did this PR introduce the regression?
(if NVPTX already had this bug, I agree with justin that we should fix it in a separate PR; i just want to rule out a PR from introducing that regression)
There was a problem hiding this comment.
@akshayrdeodhar atom.add.bf16 should already not have been emitted, because we set the smallest atomic width supported to 32-bit, so that AtomicExpand always uses a cas-loop for anything smaller.
atom.add.noftz.f16 and atom.add.noftz.bf16 are in atomicrmw-sm90.ll before this PR. I think that answers your question?
There was a problem hiding this comment.
I agree this seems unrelated to the current patch.
FWIW I think that there is some merit in emitting supported PTX instructions even when they are not 1:1 with current HW. The PTX instruction may enable early ptxas optimizations and the expansion may be more efficient if done by the final backend.
There was a problem hiding this comment.
Was NVPTX accidentally opting out of that here for bf16 and float?
This is correct. The current PR did not introduce a regression. NVPTX already has the bug. This is because shouldExpandAtomicRMWInIR, not setMinCmpXchgSizeInBits dictates emulation for low-bitwidth atomicrmw.
I can create a separate PR for fixing this.
|
Re: semantics change. Given that we can't decide for the users what they want or need, and I can see both scenarios being desired in different use cases, it may justify a flag that would let LLVM generate atomics on FP regardless of FTX mode. By default we should do things by the book, but if the user really wants fast atomic ops and does not care about FTZ, they can explicitly override it. |
|
We could add that, but if someone has to pass an additional flag it might be almost as easy for them to change to inline asm? If you're worried about the slowdown, perhaps you could use your Powers Of Google Builds to see if anyone actually even hits the case where we're making an atom.add into a CAS loop and that actually slows them down? |
Not necessarily. E.g. the code is JIT-ted at runtime, and it's not quite clear who/where/how does it. It's basically an escape hatch to buy users some time to find the place where we may want to use inline asm. E.g. if the change lands in the production compiler, and the slowdown is discovered a week or two later, by which point it will be hard to roll back the change. It's not that I'm worried about the slowdown per se. I'm worried about my options if we hit the slowdown. If the failure mode would be an obvious crash, I'd be fine without any extra knobs, but performance impact discovery tends to be delayed by quite a bit, so my concern is about having a way out of a potential jam. If things go smoothly, we can remove the flag after a while (famous last words, I know, but one can hope). |
| AI->getOperation() == AtomicRMWInst::BinOp::FAdd) { | ||
| const bool FTZ = | ||
| AI->getFunction()->getDenormalMode(APFloat::IEEEsingle()).Output == | ||
| DenormalMode::PreserveSign; |
There was a problem hiding this comment.
Should this be != DenormalMode::IEEE? Dynamic can be interpreted as we want. I'm unsure if PTX does PreserveSign or PositiveZero. @gonzalobg can you weigh in?
There was a problem hiding this comment.
I do not think "dynamic" means "anything goes". In any case I would say that a change here should be done as a separate PR.
There was a problem hiding this comment.
I do not think "dynamic" means "anything goes".
Ah, OK. If PTX provides PreserveSign, then I'm happy with the current change
There was a problem hiding this comment.
Current implementation of atom.add.f32 on global memory flushes subnormal inputs and results to sign-preserving zero; whereas atom.add.f32 on shared memory supports subnormal inputs and results and doesn’t flush them to zero.
Makes sense to me.
The PTX spec does not mention anything about FTZ for f64, but support for f64 adds is available.
| .f64 | Not supported | .add | Not supported |
|---|
|
Thank you for all the discussion. Right now my inclination is to add to the existing automated tests (trying not to make it a 3x explosion, we'll see) and then we can decide on #200732 at our liesure. I'll also add the chicken switch @Artem-B is asking for. I'll make it so that if you flip this switch, nothing should go slower than it does today. But before I do that, I'd like an answer to this thread. #200732 (comment) |
5b737bc to
e885820
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
e885820 to
5a1d79d
Compare
Previously we unconditionally lowered LLVM atomicrmw fadd to PTX atom.add. This is incorrect, because it ignores the FTZ behavior of the LLVM and PTX instructions: - atom.add.f32 on global memory flushes denormals - atom.add.f32 on shared memory does not flush denormals - atom.add.f16 and atom.add.bf16 never flush denormals We now lower to atom.add only when the function's denormal (FTZ) mode matches that of the instruction; otherwise we fall back to a CAS loop. Add -nvptx-allow-ftz-atomics as an escape hatch to use atom.add regardless of the FTZ mismatch, trading correct denormal handling for performance. Tests are regenerated via the two-step process (atomicrmw.py then update_llc_test_checks.py); the generator now emits each test with and without -denormal-fp-math-f32=preserve-sign. The new flag is covered by atomicrmw-allow-ftz-atomics.ll. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
OK, I've added the new flag and the new tests. |
5a1d79d to
76b368f
Compare
| """; RUN: llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} | FileCheck %s --check-prefixes=SM${sm},SM${sm}-NOFTZ | ||
| ; RUN: llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} -denormal-fp-math-f32=preserve-sign | FileCheck %s --check-prefixes=SM${sm},SM${sm}-FTZ | ||
| ; RUN: %if ptxas-sm_${sm} && ptxas-isa-${ptxfp} %{ llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} | %ptxas-verify -arch=sm_${sm} %} | ||
| ; RUN: %if ptxas-sm_${sm} && ptxas-isa-${ptxfp} %{ llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} -denormal-fp-math-f32=preserve-sign | %ptxas-verify -arch=sm_${sm} %} |
There was a problem hiding this comment.
What are the semantics of clang's -denormal-fp-math-f32=preserve-sign?
Does it require the:
- inputs of an FP operation to be flushed to zero,
- outputs of an FP operation to be flushed to zero, or
- both
?
If the answer is not "both" then the codegen above may be incorrect (min.ftz.f32flushes both inputs and outputs to sign-preserved-zero).
There was a problem hiding this comment.
What are the semantics of clang's -denormal-fp-math-f32=preserve-sign?
clang's -fdenormal-fp-math-f32 flag sets the LLVM attribute denormal_fpenv(float: preservesign) on every function in the module.
The semantics of denormal_fpenv(float: preservesign) are described in the LLVM langref.
If the answer is not "both" then the codegen above may be incorrect (min.ftz.f32 flushes both inputs and outputs to sign-preserved-zero).
My reading of the langref is that it requires both. Please let me know if your reading is different!
There was a problem hiding this comment.
My reading is the same, thanks!
| "nvptx-allow-ftz-atomics", cl::Hidden, | ||
| cl::desc("NVPTX Specific: Lower atomicrmw fadd to atom.add even when its " | ||
| "FTZ behavior does not match the function's denormal mode."), | ||
| cl::init(false)); |
There was a problem hiding this comment.
What happens if an application:
- compiles LLVM Module A without this flag (noftz atomics),
- compiles LLVM Module B with this flag (ftz atomics),
- B calls A::fn, and
- LTO is enabled so that A::fn is inlined into B
?
An alternative is for NVPTX to correctly compile the IR that it is given, and if a frontend wants some of the atomics it generates to use FTZ, that frontend just sets FTZ for those atomics. What are the downsides of this approach?
(note: in the short term adding this escape hatch may be fine, but I want to agree on what the right long term solution for this should be first)
There was a problem hiding this comment.
if a frontend wants some of the atomics it generates to use FTZ, that frontend just sets FTZ for those atomics. What are the downsides of this approach?
I don't believe LLVM IR allows you to set FTZ on individual instructions. So I would say "this is not possible with LLVM IR as it exists today" is the downside to the approach. :)
| if (AI->isFloatingPointOperation() && | ||
| AI->getOperation() == AtomicRMWInst::BinOp::FAdd) { | ||
| const bool FTZ = | ||
| AI->getFunction()->getDenormalMode(APFloat::IEEEsingle()).Output == |
There was a problem hiding this comment.
Do we also need to check the denormal mode of the inputs?
There was a problem hiding this comment.
I'm preserving the old behavior here, but I agree in principle we could check both and emit a lowering error (probably not here, but somewhere) if the denormal mode is anything other than ieee|ieee or preservesign|preservesign.
| """; RUN: llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} | FileCheck %s --check-prefixes=SM${sm},SM${sm}-NOFTZ | ||
| ; RUN: llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} -denormal-fp-math-f32=preserve-sign | FileCheck %s --check-prefixes=SM${sm},SM${sm}-FTZ | ||
| ; RUN: %if ptxas-sm_${sm} && ptxas-isa-${ptxfp} %{ llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} | %ptxas-verify -arch=sm_${sm} %} | ||
| ; RUN: %if ptxas-sm_${sm} && ptxas-isa-${ptxfp} %{ llc < %s -march=nvptx64 -mcpu=sm_${sm} -mattr=+ptx${ptx} -denormal-fp-math-f32=preserve-sign | %ptxas-verify -arch=sm_${sm} %} |
There was a problem hiding this comment.
What attributes does clang attach to functions in the module when -denormal-fp-math-f32=preserve-sign is set?
denormal_fpenv(preservesign), ordenormal_fpenv(preservesign | preservesign),- or
denormal_fpenv( somethingelse | preservesign)
?
There was a problem hiding this comment.
- denormal_fpenv(preservesign)
- denormal_fpenv(preservesign | preservesign)
- denormal_fpenv( somethingelse | preservesign)
The first two are the same, per the langref.
Anyway it does the first one. :) https://cuda.godbolt.org/z/TKzfdGPqb
|
I'm going to merge this since the discussion at this point seems to be about how the NVPTX backend handles FTZ in general, not about the correctness of this patch. Thank you for the reviews! |
|
As it happens, the change did cause performance regressions, and the escape hatch option proved to be useful. The investigation summary says: Take it with a grain of salt. I'm working on obtaining a reproducer. |
|
Reproducer: https://godbolt.org/z/nPGsEnMzf It appears that FP atomics on global memory always emit CAS, and the atomics on shared memory always emit atom.add, regardless of the denormal-fp-math mode. If I understand this comment in the patch, this should not be the case, and generated instructions for shared and global should flip/flop between CAS and direct op, depending on FTZ/IEEE mode in opposite directions for global/shared F32 atomics. llvm-project/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp Lines 7461 to 7466 in 39bec9c |
|
It looks like I guess it's a WAI as LLVM only guarantees forward compatibility for the bitcode only. I'm just surprised, because autoupgrade of intrinsics is working fine on the textual IR. OK. I think we're done here. There seems to be nothing wrong with LLVM and the issue is indeed completely on the user end. |
|
Makes me wonder if we want to have Otherwise it's hard to generate consistently fast FP atomics, in particular when the same function may be mixing them in global and shared AS, where no single answer would give you fast atomics for both spaces. |
I...hope you're not using fastmath for anything you care about, Art? I have found so many bugs, I kind of stopped filing them. fastmath is the wild west. Certainly having a denormal_fpenv(fast) mode sounds reasonable to me. |
|
I do not have strong opinions on fast-math, nor am I sure that FTZ/speed trade-off needs to be tied to it. My goal is more pragmatic. I have customers that care more about performance than about FTZ (or OK with specific FTZ modes that can be lowered to fast atomics), and that works for them. Before the patch things "happend to work" for them. After the patch they no longer do, and we have no way (other than the escape hatch option) to guarantee that we'll generate such code. Whichever FTZ mode we chose for a function, it will produce the slow code path for either global or shared, and always slow code for generic pointers. That's not a great situation. The disconnect is between the per-function denorm_fpenv attribute, and per-instruction behavior. Always slow atomics on generic pointers make it even more fun. We need a canonical way to express this "fast-atomics, please" preference. |
|
Totally agree with you. It just sounded like you were contemplating turning on fastmath for your customers, which I am wary of. Note however that on sm80/90/100, an atomic shared fp op is lowered to a CAS loop in SASS, https://godbolt.org/z/M8E9h1rhP. My expectation is that an LLVM CAS loop is not significantly worse than a SASS CAS loop, and it sounds like the nvidia engineers on this thread agree. If that's true, the regression here is only for atomics on generic addresses. Which XLA probably doesn't even emit! |
Whatever it was, we managed to hit it. :-) It's hard to dig out what exactly caused the regression in this specific case, but it was reliably measurable, bisected to this patch, and fixed by the escape hatch flag. Eyeballing the IR dump from a smaller example in the same build (we do not see regression there, but XLA's code generation should run the same way), we do not seem to set any explicit FTZ controls, so it's possible that it were the global atomics that switched to CAS. It's correct from the correctness standpoint, as we default to IEEE mode now and want to avoid flushing, but contrary to user's expectations of fast atomics. To think of it, even |
Yikes, maybe I ended up fixing a bug for you all. :) Maybe you want per-op denormal control then. But I suspect there is a good reason that denormal control is per-function in LLVM... |
.. and then I start thinking of inlining those functions, and what to do when they want different things. I'm sure this rabbithole goes deeper. My cop-out is to keep things working by the book by default (your patch should stay), but give users controlled access to the armory of footguns, if that's the tool they deliberately want to use for whatever reasons. The escape hatch option works, but it's too indiscriminate. I wonder if I should just tell the users to use an intrinsic that lowers to the atomic op, and call it a day. It may be a bit of a pain on the source level, but considering that they actively want to generate a "just so" instruction which may violate compiler's correctness contract, it may be a reasonable price to pay. |
|
News from the "what could possibly go wrong" department. Atomics that get lowered as a CAS loop, introduce thread divergence. Straightworward code that does an atomic op and returns the result now needs __syncthreads() to make sure that everyone doing atomic op on the same address finished the job. Otherwise it may end up running ahead of other threads that are still updating that location one at a time. It's not a new problem, but now that we use CAS way more often makes atomic use rather problematic. @gonzalobg WDYT? |
Can you give an example? I don't see how this is a bug, or even how it's observable. (Unless ptxas has a bug.) |
This is a follow-up to llvm#200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
This is a follow-up to llvm#200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
This is a follow-up to llvm#200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
This is a follow-up to #200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
…ue (#206154) This is a follow-up to llvm/llvm-project#200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
…ue (#206154) This is a follow-up to llvm/llvm-project#200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior.
This is a follow-up to #200732. Setting the default value of nvptx-allow-ftz-atomics to true preserves the status quo. It allows the NVPTX backend to emit native floating-point atomic add instructions by default, rather than falling back to the CAS loop required for strict FTZ compliance. Correct FTZ handling forces a lot of existing code to emit CAS loops. This introduces thread divergence, breaking code that relies on atomic operations being non-divergent (for example, warp-wide atomic additions followed by a load of the reduced value without explicit synchronization). While FTZ-correct atomic behavior is desirable as a long-term default, introducing thread divergence into a common operation that previously "happened to work" without synchronization is highly disruptive. Defaulting to true makes the FTZ-correct behavior opt-in for now, allowing users time to audit and adapt their code. Test cases are updated to verify both enabled and disabled states of the option, alongside the new default behavior. (cherry picked from commit 2b6e9d2)
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using !atomic.ignore.denormal.mode metadata (#217585, #217586), which is supported in Clang, MLIR, and XLA. Users who need the previous behavior globally can still pass -nvptx-allow-ftz-atomics=true. Update test cases to verify the new default disallow behavior.
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using !atomic.ignore.denormal.mode metadata (#217585, #217586), which is supported in Clang, MLIR, and XLA. Users who need the previous behavior globally can still pass -nvptx-allow-ftz-atomics=true. Update test cases to verify the new default disallow behavior.
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using !atomic.ignore.denormal.mode metadata (#217585, #217586), which is supported in Clang, MLIR, and XLA. Users who need the previous behavior globally can still pass -nvptx-allow-ftz-atomics=true. Update test cases to verify the new default disallow behavior.
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using !atomic.ignore.denormal.mode metadata (#217585, #217586), which is supported in Clang, MLIR, and XLA. Users who need the previous behavior globally can still pass -nvptx-allow-ftz-atomics=true. Update test cases to verify the new default disallow behavior.
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using !atomic.ignore.denormal.mode metadata (#217585, #217586), which is supported in Clang, MLIR, and XLA. Users who need the previous behavior globally can still pass -nvptx-allow-ftz-atomics=true. Update test cases to verify the new default disallow behavior.
This restores the strict FTZ default originally introduced in #200732 and temporarily relaxed in #206154. Frontends can now explicitly opt in to native FTZ-mismatched atomics on a per-instruction basis using `!atomic.ignore.denormal.mode` metadata (#217585, #217586, #217587, #217588, #232725, #232726), which is supported in Clang, Flang, MLIR, and XLA. Users who need the previous behavior globally can still pass `-nvptx-allow-ftz-atomics=true`. Update test cases to verify the new default disallow behavior.
Previously we unconditionally lowered LLVM atomicrmw fadd to PTX
atom.add. This is incorrect, because it ignores the FTZ behavior of the
LLVM and PTX instructions.
Co-Authored-By: Claude Opus 4.8 (1M context) noreply@anthropic.com