[NVPTX] Enable AtomicExpandPass for NVPTX

[tianshilei1992]

Summary

This patch enables AtomicExpandPass for NVPTX.

Depend on D125652.

[tianshilei1992]

llvm/lib/CodeGen/AtomicExpandPass.cpp

257

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 ↗

(On Diff #429584)

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 ↗

(On Diff #429584)

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 ↗

(On Diff #429584)

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 ↗

(On Diff #429842)

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 ↗

(On Diff #429842)

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 ↗

(On Diff #429842)

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 ↗

(On Diff #429842)

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 ↗

(On Diff #429842)

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/CodeGen/AtomicExpandPass.cpp

if (TLI->shouldInsertFencesForAtomic(I)) {
       }
 
       if (FenceOrdering != AtomicOrdering::Monotonic) {
         MadeChange |= bracketInstWithFences(I, FenceOrdering);
       }
     }
 
     if (LI) {
+      if (TLI->shouldExpandAtomicLoadInIR(LI) !=
+          TargetLoweringBase::AtomicExpansionKind::None) {
         if (LI->getType()->isFloatingPointTy()) {
           // TODO: add a TLI hook to control this so that each target can
           // convert to lowering the original type one at a time.
           LI = convertAtomicLoadToIntegerType(LI);
           assert(LI->getType()->isIntegerTy() && "invariant broken");
           MadeChange = true;
         }
 
         MadeChange |= tryExpandAtomicLoad(LI);
+      }
     } else if (SI) {
+      if (TLI->shouldExpandAtomicStoreInIR(SI) !=
+          TargetLoweringBase::AtomicExpansionKind::None) {
         if (SI->getValueOperand()->getType()->isFloatingPointTy()) {
           // TODO: add a TLI hook to control this so that each target can
           // convert to lowering the original type one at a time.
           SI = convertAtomicStoreToIntegerType(SI);
           assert(SI->getValueOperand()->getType()->isIntegerTy() &&
                  "invariant broken");
           MadeChange = true;
         }
 
         if (tryExpandAtomicStore(SI))
           MadeChange = true;
+      }
     } else if (RMWI) {
       // There are two different ways of expanding RMW instructions:
       // - into a load if it is idempotent
       // - into a Cmpxchg/LL-SC loop otherwise
       // we try them in that order.
 
       if (isIdempotentRMW(RMWI) && simplifyIdempotentRMW(RMWI)) {
         MadeChange = true;

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
   //