[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

70

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 ↗

(On Diff #340182)

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

32 ↗

(On Diff #340182)

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/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 "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 &function) override;
+};
+} // namespace
+
+bool NVPTXAtomicLower::runOnFunction(Function &F) {
+  SmallVector<AtomicRMWInst *, 3> V;
+  for (auto &BB : F) {
+    for (auto &I : BB) {
+      if (auto N = dyn_cast<AtomicRMWInst>(&I)) {
+        if (N->getPointerAddressSpace() == ADDRESS_SPACE_LOCAL &&
+            N->getOrdering() == AtomicOrdering::Monotonic) {
+          V.push_back(N);
+        }
+      }
+    }
+  }
+  bool Changed = false;
+  for (auto I : V) {
+    if (I->getOperation() == llvm::AtomicRMWInst::BinOp::FAdd) {
+      IRBuilder<> B(I);
+      Changed = true;
+      auto L = B.CreateLoad(I->getType(), I->getPointerOperand());
+      L->setAtomic(I->getOrdering(), I->getSyncScopeID());
+      auto SI = B.CreateStore(B.CreateFAdd(L, I->getValOperand()),
+                              I->getPointerOperand());
+      SI->setAtomic(I->getOrdering(), I->getSyncScopeID());
+      I->replaceAllUsesWith(L);
+      I->eraseFromParent();
+    }
+  }
+  return Changed;
+}
+
+char NVPTXAtomicLower::ID = 0;
+
+namespace llvm {
+void 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"

 }
 
 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());