[llvm] [NVPTX] Respect FTZ flag when lowering atomicrmw fadd. (PR #200732)

[Justin Lebar via llvm-commits]

================
@@ -7467,22 +7467,52 @@ 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()) {
----------------
jlebar wrote:

> With the current change, an operation in generic will result in hitting the AI->isFloatingPointOperation() check below, expanding to a cmpxchg loop.

Correct.

> Was the old behavior (do not expand) for the generic space incorrect?

According to the PTX docs, atom.add.f32 on global memory flushes to zero, whereas atom.add.f32 on shmem does not.  Therefore if a pointer is in generic memory, we don't know whether or not atom.add will flush to zero.  So under the current patch we can never use it.

https://github.com/llvm/llvm-project/pull/200732