[NVPTX] Enable lowering of atomics on local memory

[wsmoses]

Summary

LLVM does not have valid assembly backends for atomicrmw on local memory. However, as this memory is thread local, we should be able to lower this to the relevant load/store.

[tra]

I think the right place to deal with this is during the lowering. This pass just works around the actual problem.

llvm/lib/Target/NVPTX/NVPTXAtomicLower.cpp

80

The description needs to be updated.

[wsmoses]

Address comment regarding description

[jdoerfert]

No test.

llvm/lib/Target/NVPTX/NVPTXAtomicLower.cpp

71

V is not a descriptive name.
auto* where appropriate.
if the type is not clean, don't use auto.
Why only for monotonic orderings?
Why only FAdd?
Don't we have an existing way to "deconstruct" atomics intriniscs into loads/stores?

[wsmoses]

Add test, rename variables, and use existing atomic lower code

[jdoerfert]

Only minor nits from me, @tra, you should review this I suppose.

llvm/lib/Target/NVPTX/NVPTXAtomicLower.cpp

46

55

Nit: for (Instructions &I : instructions(F))
I would omit the braces, also below.

[tra]

Few nits. LGTM overall.

llvm/include/llvm/Transforms/Scalar/LowerAtomic.h

17

We could probably get by with a class AtomicRMWInst;. We don't need the whole header just for the pointer type.

31

Nit: function names should be camel case, and start with a lower case letter (e.g. openFile() or isFoo())., according to LLVM style guide: https://llvm.org/docs/CodingStandards.html#name-types-functions-variables-and-enumerators-properly

llvm/lib/Target/NVPTX/NVPTXAtomicLower.cpp

65–67

llvm::initializeNVPTXAtomicLowerPass(PassRegistry &); would do.

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

7–12 ↗

(On Diff #340182)

The file should be renamed to atomic-lower-local.ll and there should be a comment describing that the test ensures that we're able to lower atomics in local AS in principle and that in case of NVPTX we happen to lower them to regular load/stores. Without the explanation, the CHECK lines by themselves are somewhat puzzling as there is nothing about atomics in the IR they check for.

[wsmoses]

Address comments

llvm/include/llvm/Transforms/Scalar/LowerAtomic.h

 //===- LowerAtomic.cpp - Lower atomic intrinsics ----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file
 // This pass lowers atomic intrinsics to non-atomic form for use in a known
 // non-preemptible environment.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_SCALAR_LOWERATOMIC_H
 #define LLVM_TRANSFORMS_SCALAR_LOWERATOMIC_H
 
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// A pass that lowers atomic intrinsic into non-atomic intrinsics.
 class LowerAtomicPass : public PassInfoMixin<LowerAtomicPass> {
 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
   static bool isRequired() { return true; }
 };
+
+class AtomicRMWInst;
+/// Convert the given RMWI into primitive load and stores,
+/// assuming that doing so is legal. Return true if the lowering
+/// succeeds.
+bool lowerAtomicRMWInst(AtomicRMWInst *RMWI);
 }
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOWERATOMIC_H

llvm/lib/Target/NVPTX/CMakeLists.txt

 add_llvm_component_group(NVPTX)
 
 set(LLVM_TARGET_DEFINITIONS NVPTX.td)
 
 tablegen(LLVM NVPTXGenAsmWriter.inc -gen-asm-writer)
 tablegen(LLVM NVPTXGenDAGISel.inc -gen-dag-isel)
 tablegen(LLVM NVPTXGenInstrInfo.inc -gen-instr-info)
 tablegen(LLVM NVPTXGenRegisterInfo.inc -gen-register-info)
 tablegen(LLVM NVPTXGenSubtargetInfo.inc -gen-subtarget)
 
 add_public_tablegen_target(NVPTXCommonTableGen)
 
 set(NVPTXCodeGen_sources
   NVPTXAllocaHoisting.cpp
+  NVPTXAtomicLower.cpp
   NVPTXAsmPrinter.cpp
   NVPTXAssignValidGlobalNames.cpp
   NVPTXFrameLowering.cpp
   NVPTXGenericToNVVM.cpp
   NVPTXISelDAGToDAG.cpp
   NVPTXISelLowering.cpp
   NVPTXImageOptimizer.cpp
   NVPTXInstrInfo.cpp

llvm/lib/Target/NVPTX/NVPTXAtomicLower.h

+//===-- NVPTXAtomicLower.h - Lower atomics of local memory ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//  Lower atomics of local memory to simple load/stores
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_NVPTX_NVPTXATOMICLOWER_H
+#define LLVM_LIB_TARGET_NVPTX_NVPTXATOMICLOWER_H
+
+namespace llvm {
+class FunctionPass;
+
+extern FunctionPass *createNVPTXAtomicLowerPass();
+} // end namespace llvm
+
+#endif

llvm/lib/Target/NVPTX/NVPTXAtomicLower.cpp

+//===-- NVPTXAtomicLower.cpp - Lower atomics of local memory ----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//  Lower atomics of local memory to simple load/stores
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTXAtomicLower.h"
+#include "llvm/CodeGen/StackProtector.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/Transforms/Scalar/LowerAtomic.h"
+
+#include "MCTargetDesc/NVPTXBaseInfo.h"
+using namespace llvm;
+
+namespace {
+// Hoisting the alloca instructions in the non-entry blocks to the entry
+// block.
+class NVPTXAtomicLower : public FunctionPass {
+public:
+  static char ID; // Pass ID
+  NVPTXAtomicLower() : FunctionPass(ID) {}
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+  }
+
+  StringRef getPassName() const override {
+    return "NVPTX lower atomics of local memory";
+  }
+
+  bool runOnFunction(Function &F) override;
+};
+} // namespace
+
+bool NVPTXAtomicLower::runOnFunction(Function &F) {
+  SmallVector<AtomicRMWInst *> LocalMemoryAtomics;
+  for (Instruction &I : instructions(F))
+    if (AtomicRMWInst *RMWI = dyn_cast<AtomicRMWInst>(&I))
+      if (RMWI->getPointerAddressSpace() == ADDRESS_SPACE_LOCAL)
+        LocalMemoryAtomics.push_back(RMWI);
+
+  bool Changed = false;
+  for (AtomicRMWInst *RMWI : LocalMemoryAtomics)
+    Changed |= lowerAtomicRMWInst(RMWI);
+  return Changed;
+}
+
+char NVPTXAtomicLower::ID = 0;
+
+void llvm::initializeNVPTXAtomicLowerPass(PassRegistry &);
+
+INITIALIZE_PASS(NVPTXAtomicLower, "nvptx-atomic-lower",
+                "Lower atomics of local memory to simple load/stores", false,
+                false)
+
+FunctionPass *llvm::createNVPTXAtomicLowerPass() {
+  return new NVPTXAtomicLower();
+}

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

 //===-- NVPTXTargetMachine.cpp - Define TargetMachine for NVPTX -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Top-level implementation for the NVPTX target.
 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTXTargetMachine.h"
 #include "NVPTX.h"
 #include "NVPTXAllocaHoisting.h"
+#include "NVPTXAtomicLower.h"
 #include "NVPTXLowerAggrCopies.h"
 #include "NVPTXTargetObjectFile.h"
 #include "NVPTXTargetTransformInfo.h"
 #include "TargetInfo/NVPTXTargetInfo.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/Passes.h"
static cl::opt<bool> UseShortPointersOpt(
     cl::init(false), cl::Hidden);
 
 namespace llvm {
 
 void initializeNVVMIntrRangePass(PassRegistry&);
 void initializeNVVMReflectPass(PassRegistry&);
 void initializeGenericToNVVMPass(PassRegistry&);
 void initializeNVPTXAllocaHoistingPass(PassRegistry &);
+void initializeNVPTXAtomicLowerPass(PassRegistry &);
 void initializeNVPTXAssignValidGlobalNamesPass(PassRegistry&);
 void initializeNVPTXLowerAggrCopiesPass(PassRegistry &);
 void initializeNVPTXLowerArgsPass(PassRegistry &);
 void initializeNVPTXLowerAllocaPass(PassRegistry &);
 void initializeNVPTXProxyRegErasurePass(PassRegistry &);
 
 } // end namespace llvm
 
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   // Register the target.
   RegisterTargetMachine<NVPTXTargetMachine32> X(getTheNVPTXTarget32());
   RegisterTargetMachine<NVPTXTargetMachine64> Y(getTheNVPTXTarget64());
 
   // FIXME: This pass is really intended to be invoked during IR optimization,
   // but it's very NVPTX-specific.
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   initializeNVVMReflectPass(PR);
   initializeNVVMIntrRangePass(PR);
   initializeGenericToNVVMPass(PR);
   initializeNVPTXAllocaHoistingPass(PR);
   initializeNVPTXAssignValidGlobalNamesPass(PR);
+  initializeNVPTXAtomicLowerPass(PR);
   initializeNVPTXLowerArgsPass(PR);
   initializeNVPTXLowerAllocaPass(PR);
   initializeNVPTXLowerAggrCopiesPass(PR);
   initializeNVPTXProxyRegErasurePass(PR);
 }
 
 static std::string computeDataLayout(bool is64Bit, bool UseShortPointers) {
   std::string Ret = "e";

 }
 
 void NVPTXPassConfig::addAddressSpaceInferencePasses() {
   // NVPTXLowerArgs emits alloca for byval parameters which can often
   // be eliminated by SROA.
   addPass(createSROAPass());
   addPass(createNVPTXLowerAllocaPass());
   addPass(createInferAddressSpacesPass());
+  addPass(createNVPTXAtomicLowerPass());
 }
 
 void NVPTXPassConfig::addStraightLineScalarOptimizationPasses() {
   addPass(createSeparateConstOffsetFromGEPPass());
   addPass(createSpeculativeExecutionPass());
   // ReassociateGEPs exposes more opportunites for SLSR. See
   // the example in reassociate-geps-and-slsr.ll.
   addPass(createStraightLineStrengthReducePass());

llvm/lib/Transforms/Scalar/LowerAtomic.cpp

static bool LowerAtomicCmpXchgInst(AtomicCmpXchgInst *CXI) {
   Res = Builder.CreateInsertValue(UndefValue::get(CXI->getType()), Orig, 0);
   Res = Builder.CreateInsertValue(Res, Equal, 1);
 
   CXI->replaceAllUsesWith(Res);
   CXI->eraseFromParent();
   return true;
 }
 
-static bool LowerAtomicRMWInst(AtomicRMWInst *RMWI) {
+bool llvm::lowerAtomicRMWInst(AtomicRMWInst *RMWI) {
   IRBuilder<> Builder(RMWI);
   Value *Ptr = RMWI->getPointerOperand();
   Value *Val = RMWI->getValOperand();
 
   LoadInst *Orig = Builder.CreateLoad(Val->getType(), Ptr);
   Value *Res = nullptr;
 
   switch (RMWI->getOperation()) {

 static bool runOnBasicBlock(BasicBlock &BB) {
   bool Changed = false;
   for (Instruction &Inst : make_early_inc_range(BB)) {
     if (FenceInst *FI = dyn_cast<FenceInst>(&Inst))
       Changed |= LowerFenceInst(FI);
     else if (AtomicCmpXchgInst *CXI = dyn_cast<AtomicCmpXchgInst>(&Inst))
       Changed |= LowerAtomicCmpXchgInst(CXI);
     else if (AtomicRMWInst *RMWI = dyn_cast<AtomicRMWInst>(&Inst))
-      Changed |= LowerAtomicRMWInst(RMWI);
+      Changed |= lowerAtomicRMWInst(RMWI);
     else if (LoadInst *LI = dyn_cast<LoadInst>(&Inst)) {
       if (LI->isAtomic())
         LowerLoadInst(LI);
     } else if (StoreInst *SI = dyn_cast<StoreInst>(&Inst)) {
       if (SI->isAtomic())
         LowerStoreInst(SI);
     }
   }

llvm/test/CodeGen/NVPTX/atomic-lower-local.ll

+; RUN: opt < %s -S -nvptx-atomic-lower | FileCheck %s
+
+; This test ensures that there is a legal way for ptx to lower atomics
+; on local memory. Here, we demonstrate this by lowering them to simple
+; load and stores.
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+define double @kernel(double addrspace(5)* %ptr, double %val) {
+  %res = atomicrmw fadd double addrspace(5)* %ptr, double %val monotonic, align 8
+  ret double %res
+; CHECK:   %1 = load double, double addrspace(5)* %ptr, align 8
+; CHECK-NEXT:   %2 = fadd double %1, %val
+; CHECK-NEXT:   store double %2, double addrspace(5)* %ptr, align 8
+; CHECK-NEXT:   ret double %1
+}
+