[NVPTX] Enable AtomicExpandPass for NVPTX

[tianshilei1992]

Summary

This patch enables AtomicExpandPass for NVPTX.

Depend on D125652.

[tianshilei1992]

llvm/lib/CodeGen/AtomicExpandPass.cpp

257 ↗

(On Diff #429550)

Some targets, such as NVPTX, supports atomic load and store of floating-point variables. We don't need to convert it to integer.

[tianshilei1992]

rebase

[tianshilei1992]

add test

[tra]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5159

According to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-atom

64-bit atom.{and,or,xor,min,max} require sm_32 or higher.

This must be conditional on the GPU variant we're compiling for, similar to how we handle f64 above.

llvm/test/CodeGen/NVPTX/atomic-expand.ll.ll

7

Your code also adds support for atomics on integers, those should be verified, too.

In general we do want to have tests for all types native to PTX.

[tianshilei1992]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5159

Thanks for pointing it out. Will do.

I got a question actually. What's the minimum SM version do we currently support? I did find from one of your patch (https://reviews.llvm.org/D24943) that it had the function hasAtomAddF32, but now I can't find it. Was it removed because we support a minimum SM version, like SM20?

[tra]

What's the minimum SM version do we currently support? I did find from one of your patch (https://reviews.llvm.org/D24943) that it had the function hasAtomAddF32, but now I can't find it. Was it removed because we support a minimum SM version, like SM20?

SM20 is the current minimum. Note that NVIDIA has already stopped supporting pre-SM35 GPUs and we should probably start considering removing support for sm_2x in LLVM, too.

[tianshilei1992]

fix comments

[tra]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5160

Can we handle types other than 32 and 64-bit? If yes, we need tests at least for some of them (i8/i16/f16/i128). If not, this should probably be llvm_unreachable().

llvm/test/CodeGen/NVPTX/atomic-expand.ll

1 ↗

(On Diff #429842)

Nit: You can specify multiple labels as one option --check-prefixes=A,B,C

I'd also rename the labels.

CHECK->COMMON or ALL -- makes the intent clear.

CHECK-SM30 -> SM30 -- reduces unnecessary clutter and makes it a bit easier to spot the differences between SM30 and SM60.

[tianshilei1992]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5160

This is quite interesting. I don't think CUDA can generate that kind of code, but OpenMP can. However, for other types, like i16, with the atomic expand, it's just converted to AtomicCmpSwap of type i16, and then the isel crashed. For now, I guess the best way is to use llvm_unreachable for other cases. In the long term, we may want to support the expansion of types with different bitwidth, but I don't have a clear idea how that would work. Otherwise, we will have to tell the front end that some types are not supported, which kind of breaks the assumption that front end can emit target agnostic code.

[tra]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5160

I don't think CUDA can generate that kind of code,

Why do you believe that's the case? I would assume that if a C++ compilation can, so would CUDA, though we may be sort of lucky and may currently not have appropriate GPU-side overloads for std::atomic functions. We do want to make it work, so it's likely to get fixed sooner or later.

It may still be a good idea to add these extra test cases to the test, too, possibly commented out with a comment explaining what's going on.

[tianshilei1992]

fix comments

[tianshilei1992]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5160

I added tests for i8 and i16. Another types are not covered because based on https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#type-system they are not supported at the IR level.

[tra]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

5160

they are not supported [by NVVM] at the IR level.

NVVM != LLVM, so it's not a particularly strong reason.

LLVM does have limited support for i128 and we should cover it, even if the answer is "it does not work (yet?)".

[efriedma]

setMinCmpXchgSizeInBits(32) should fix the i8/i16 testcases, I think.

[tianshilei1992]

add i128

[tianshilei1992]

use setMinCmpXchgSizeInBits(32)

[tianshilei1992]

In D125639#3523733, @efriedma wrote:

setMinCmpXchgSizeInBits(32) should fix the i8/i16 testcases, I think.

Yes, that works. Thanks for the info.

[tianshilei1992]

fix return

[tianshilei1992]

fix tests

llvm/lib/Target/NVPTX/NVPTXISelLowering.h

public:
   bool enableAggressiveFMAFusion(EVT VT) const override { return true; }
 
   // The default is to transform llvm.ctlz(x, false) (where false indicates that
   // x == 0 is not undefined behavior) into a branch that checks whether x is 0
   // and avoids calling ctlz in that case.  We have a dedicated ctlz
   // instruction, so we say that ctlz is cheap to speculate.
   bool isCheapToSpeculateCtlz() const override { return true; }
 
+  AtomicExpansionKind shouldCastAtomicLoadInIR(LoadInst *LI) const override {
+    return AtomicExpansionKind::None;
+  }
+
+  AtomicExpansionKind shouldCastAtomicStoreInIR(StoreInst *SI) const override {
+    return AtomicExpansionKind::None;
+  }
+
+  AtomicExpansionKind
+  shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const override;
+
 private:
   const NVPTXSubtarget &STI; // cache the subtarget here
   SDValue getParamSymbol(SelectionDAG &DAG, int idx, EVT) const;
 
   SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;
 

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

case ISD::LOAD:
     ReplaceLoadVector(N, DAG, Results);
     return;
   case ISD::INTRINSIC_W_CHAIN:
     ReplaceINTRINSIC_W_CHAIN(N, DAG, Results);
     return;
   }
 }
 
+NVPTXTargetLowering::AtomicExpansionKind
+NVPTXTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
+  Type *Ty = AI->getValOperand()->getType();
+  if (AI->isFloatingPointOperation()) {
+    if (AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
+      if (Ty->isFloatTy())
+        return AtomicExpansionKind::None;
+      if (Ty->isDoubleTy() && STI.hasAtomAddF64())
+        return AtomicExpansionKind::None;
+    }
+    return AtomicExpansionKind::CmpXChg;
+  }
+
+  switch (AI->getOperation()) {
+  default:
+    return AtomicExpansionKind::CmpXChg;
+  case AtomicRMWInst::BinOp::Add:
+  case AtomicRMWInst::BinOp::Sub:
+  case AtomicRMWInst::BinOp::Max:
+  case AtomicRMWInst::BinOp::Min:
+  case AtomicRMWInst::BinOp::UMax:
+  case AtomicRMWInst::BinOp::UMin:
+  case AtomicRMWInst::BinOp::Xchg:
+  case AtomicRMWInst::BinOp::And:
+  case AtomicRMWInst::BinOp::Or:
+  case AtomicRMWInst::BinOp::Xor:
+    assert(Ty->isIntegerTy());
+    switch (cast<llvm::IntegerType>(Ty)->getBitWidth()) {
+    case 32:
+      return AtomicExpansionKind::None;
+    case 64:
+      return AtomicExpansionKind::None;
+    default:
+      return AtomicExpansionKind::CmpXChg;
+    }
+  }
+
+  return AtomicExpansionKind::CmpXChg;
+}
+
 // Pin NVPTXTargetObjectFile's vtables to this file.
 NVPTXTargetObjectFile::~NVPTXTargetObjectFile() = default;
 
 MCSection *NVPTXTargetObjectFile::SelectSectionForGlobal(
     const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM) const {
   return getDataSection();
 }

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

void NVPTXPassConfig::addIRPasses() {
   // NVPTXLowerArgs is required for correctness and should be run right
   // before the address space inference passes.
   addPass(createNVPTXLowerArgsPass(&getNVPTXTargetMachine()));
   if (getOptLevel() != CodeGenOpt::None) {
     addAddressSpaceInferencePasses();
     addStraightLineScalarOptimizationPasses();
   }
 
+  addPass(createAtomicExpandPass());
+
   // === LSR and other generic IR passes ===
   TargetPassConfig::addIRPasses();
   // EarlyCSE is not always strong enough to clean up what LSR produces. For
   // example, GVN can combine
   //
   //   %0 = add %a, %b
   //   %1 = add %b, %a
   //

llvm/test/CodeGen/NVPTX/atomic-expand.ll.ll

+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefix=CHECK-SM20
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_60 | FileCheck %s --check-prefix=CHECK-SM60
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 | %ptxas-verify %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_60 | %ptxas-verify %}
+
+; CHECK-LABEL: foo
+define void @foo(ptr %0, double %1) {
+entry:
+  ; CHECK-SM20: atom.cas.b64
+  ; CHECK-SM60: atom.add.f64
+  %2 = atomicrmw fadd ptr %0, double %1 monotonic, align 8
+  ret void
+}