[RFC][Clang] Add a check if the target supports atomicrmw instruction with specific operator and type

[tianshilei1992]

Summary

With respect of atomicrmw instruction, different targets have different
backend support. The front end can always emit the intruction, but when it comes
to the backend, the compiler could crash at instruction selection because the
target doesn't support the instruction/operator at all. Currently we don't have
a way in the front end to tell if an atomicrmw with specific operator and type
is supported. This patch adds a virtual function hasAtomicrmw to the class
TargetInfo and it is expected to be implemented by all targets accordingly.
In this way, we can always try to emit atomicrmw first because it has better
performance than falling back to CAS loop. Currently at the PoC stage, I only make
it in NVPTX, but if the community is happy with this change, I'll add to all the
targets.

For the long term, I'll seek a way to expand the instruction if the target doesn't
support it, probably similar to things in llvm/lib/CodeGen/AtomicExpandPass.cpp.

[tianshilei1992]

Feel free to add more reviewers

[tianshilei1992]

clang/lib/Basic/Targets/NVPTX.cpp

311

llvm_unreachable might be more appropriate here.

[tianshilei1992]

As for the test, here is a simple code snippet:

#pragma omp begin declare target device_type(nohost)
void foo() {
  double x, y;
#pragma omp atomic update
  x += y;
}
#pragma omp end declare target

sm_35 doesn't support FAdd with double type, then it emits:

define protected void @foo() #0 {
entry:
  %x = alloca double, align 8
  %y = alloca double, align 8
  %atomic-temp = alloca double, align 8
  %0 = load double, ptr %y, align 8
  %atomic-load = load atomic i64, ptr %x monotonic, align 8
  br label %atomic_cont

atomic_cont:                                      ; preds = %atomic_cont, %entry
  %1 = phi i64 [ %atomic-load, %entry ], [ %5, %atomic_cont ]
  %2 = bitcast i64 %1 to double
  %add = fadd double %2, %0
  store double %add, ptr %atomic-temp, align 8
  %3 = load i64, ptr %atomic-temp, align 8
  %4 = cmpxchg ptr %x, i64 %1, i64 %3 monotonic monotonic, align 8
  %5 = extractvalue { i64, i1 } %4, 0
  %6 = extractvalue { i64, i1 } %4, 1
  br i1 %6, label %atomic_exit, label %atomic_cont

atomic_exit:                                      ; preds = %atomic_cont
  ret void
}

Starting from sm_60, it is supported, now we get:

define protected void @foo() #0 {
entry:
  %x = alloca double, align 8
  %y = alloca double, align 8
  %0 = load double, ptr %y, align 8
  %1 = atomicrmw fadd ptr %x, double %0 monotonic, align 8
  ret void
}

[tianshilei1992]

Do it in BE.

clang/include/clang/Basic/TargetInfo.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Frontend/OpenMP/OMPGridValues.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/Support/DataTypes.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/VersionTuple.h"
 #include <cassert>
 #include <string>
 #include <vector>
 
 namespace llvm {
public:
   virtual bool hasBuiltinAtomic(uint64_t AtomicSizeInBits,
                                 uint64_t AlignmentInBits) const {
     return AtomicSizeInBits <= AlignmentInBits &&
            AtomicSizeInBits <= getMaxAtomicInlineWidth() &&
            (AtomicSizeInBits <= getCharWidth() ||
             llvm::isPowerOf2_64(AtomicSizeInBits / getCharWidth()));
   }
 
+  /// Returns true if the given target supports atomicrmw with given operator
+  /// and type.
+  virtual bool hasAtomicrmw(llvm::AtomicRMWInst::BinOp, llvm::Type *) const {
+    llvm_unreachable("hasAtomicrmw not implemented on this target");
+  }
+
   /// Return the maximum vector alignment supported for the given target.
   unsigned getMaxVectorAlign() const { return MaxVectorAlign; }
   /// Return default simd alignment for the given target. Generally, this
   /// value is type-specific, but this alignment can be used for most of the
   /// types for the given target.
   unsigned getSimdDefaultAlign() const { return SimdDefaultAlign; }
 
   unsigned getMaxOpenCLWorkGroupSize() const { return MaxOpenCLWorkGroupSize; }

clang/lib/Basic/Targets/NVPTX.h

 
 #ifndef LLVM_CLANG_LIB_BASIC_TARGETS_NVPTX_H
 #define LLVM_CLANG_LIB_BASIC_TARGETS_NVPTX_H
 
 #include "clang/Basic/Cuda.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TargetOptions.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/Support/Compiler.h"
 
 namespace clang {
 namespace targets {
 
 static const unsigned NVPTXAddrSpaceMap[] = {
     0, // Default
     1, // opencl_global
CallingConvCheckResult checkCallingConvention(CallingConv CC) const override {
     // TODO: We should warn if you apply a non-default CC to anything other than
     // a host function.
     if (HostTarget)
       return HostTarget->checkCallingConvention(CC);
     return CCCR_Warning;
   }
 
   bool hasBitIntType() const override { return true; }
+
+  bool hasAtomicrmw(llvm::AtomicRMWInst::BinOp, llvm::Type *) const override;
 };
 } // namespace targets
 } // namespace clang
 #endif // LLVM_CLANG_LIB_BASIC_TARGETS_NVPTX_H

clang/lib/Basic/Targets/NVPTX.cpp

if (Opts.CUDAIsDevice) {
     Builder.defineMacro("__CUDA_ARCH__", CUDAArchCode);
   }
 }
 
 ArrayRef<Builtin::Info> NVPTXTargetInfo::getTargetBuiltins() const {
   return llvm::makeArrayRef(BuiltinInfo, clang::NVPTX::LastTSBuiltin -
                                              Builtin::FirstTSBuiltin);
 }
+
+bool NVPTXTargetInfo::hasAtomicrmw(llvm::AtomicRMWInst::BinOp Op,
+                                   llvm::Type *Ty) const {
+  switch (Op) {
+  default:
+    llvm_unreachable("bad binop");
+  case llvm::AtomicRMWInst::BinOp::FSub:
+  case llvm::AtomicRMWInst::BinOp::Nand:
+    return false;
+  case llvm::AtomicRMWInst::BinOp::FAdd:
+    assert(Ty->isFloatingPointTy());
+    if (Ty->isFloatTy())
+      return true;
+    if (Ty->isDoubleTy() && GPU >= CudaArch::SM_60)
+      return true;
+    return false;
+  case llvm::AtomicRMWInst::BinOp::Add:
+  case llvm::AtomicRMWInst::BinOp::Sub:
+  case llvm::AtomicRMWInst::BinOp::Max:
+  case llvm::AtomicRMWInst::BinOp::Min:
+  case llvm::AtomicRMWInst::BinOp::UMax:
+  case llvm::AtomicRMWInst::BinOp::UMin:
+  case llvm::AtomicRMWInst::BinOp::Xchg:
+  case llvm::AtomicRMWInst::BinOp::And:
+  case llvm::AtomicRMWInst::BinOp::Or:
+  case llvm::AtomicRMWInst::BinOp::Xor:
+    assert(Ty->isIntegerTy());
+    switch (cast<llvm::IntegerType>(Ty)->getBitWidth()) {
+    case 32:
+      return true;
+    case 64:
+      return true;
+    default:
+      return false;
+    }
+    return false;
+  }
+
+  return false;
+}

clang/lib/CodeGen/CGStmtOpenMP.cpp

static std::pair<bool, RValue> emitOMPAtomicRMW(CodeGenFunction &CGF, LValue X,
   case BO_MulAssign:
   case BO_DivAssign:
   case BO_RemAssign:
   case BO_ShlAssign:
   case BO_ShrAssign:
   case BO_Comma:
     llvm_unreachable("Unsupported atomic update operation");
   }
+
+  if (!Context.getTargetInfo().hasAtomicrmw(RMWOp,
+                                            X.getAddress(CGF).getElementType()))
+    return std::make_pair(false, RValue::get(nullptr));
+
   llvm::Value *UpdateVal = Update.getScalarVal();
   if (auto *IC = dyn_cast<llvm::ConstantInt>(UpdateVal)) {
     if (IsInteger)
       UpdateVal = CGF.Builder.CreateIntCast(
           IC, X.getAddress(CGF).getElementType(),
           X.getType()->hasSignedIntegerRepresentation());
     else
       UpdateVal = CGF.Builder.CreateCast(llvm::Instruction::CastOps::UIToFP, IC,