[NVPTX] Pull invariant load identification into IR pass

llvm/lib/Target/NVPTX/CMakeLists.txt

@@ -13,34 +13,35 @@ add_public_tablegen_target(NVPTXCommonTableGen)
 set(NVPTXCodeGen_sources
   NVPTXAliasAnalysis.cpp
   NVPTXAllocaHoisting.cpp
-  NVPTXAtomicLower.cpp
   NVPTXAsmPrinter.cpp
   NVPTXAssignValidGlobalNames.cpp
+  NVPTXAtomicLower.cpp
+  NVPTXCtorDtorLowering.cpp
   NVPTXForwardParams.cpp
   NVPTXFrameLowering.cpp
   NVPTXGenericToNVVM.cpp
-  NVPTXISelDAGToDAG.cpp
-  NVPTXISelLowering.cpp
   NVPTXImageOptimizer.cpp
   NVPTXInstrInfo.cpp
+  NVPTXISelDAGToDAG.cpp
+  NVPTXISelLowering.cpp
   NVPTXLowerAggrCopies.cpp
-  NVPTXLowerArgs.cpp
   NVPTXLowerAlloca.cpp
+  NVPTXLowerArgs.cpp
   NVPTXLowerUnreachable.cpp
-  NVPTXPeephole.cpp
   NVPTXMCExpr.cpp
+  NVPTXPeephole.cpp
   NVPTXPrologEpilogPass.cpp
+  NVPTXProxyRegErasure.cpp
   NVPTXRegisterInfo.cpp
   NVPTXReplaceImageHandles.cpp
   NVPTXSelectionDAGInfo.cpp
   NVPTXSubtarget.cpp
+  NVPTXTagInvariantLoads.cpp
   NVPTXTargetMachine.cpp
   NVPTXTargetTransformInfo.cpp
   NVPTXUtilities.cpp
   NVVMIntrRange.cpp
   NVVMReflect.cpp
-  NVPTXProxyRegErasure.cpp
-  NVPTXCtorDtorLowering.cpp
   )
 
 add_llvm_target(NVPTXCodeGen

llvm/lib/Target/NVPTX/NVPTX.h

@@ -51,6 +51,7 @@ FunctionPass *createNVPTXLowerArgsPass();
 FunctionPass *createNVPTXLowerAllocaPass();
 FunctionPass *createNVPTXLowerUnreachablePass(bool TrapUnreachable,
                                               bool NoTrapAfterNoreturn);
+FunctionPass *createNVPTXTagInvariantLoadsPass();
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 MachineFunctionPass *createNVPTXForwardParamsPass();
@@ -73,6 +74,7 @@ void initializeNVVMReflectPass(PassRegistry &);
 void initializeNVPTXAAWrapperPassPass(PassRegistry &);
 void initializeNVPTXExternalAAWrapperPass(PassRegistry &);
 void initializeNVPTXPeepholePass(PassRegistry &);
+void initializeNVPTXTagInvariantLoadLegacyPassPass(PassRegistry &);
 
 struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
@@ -104,6 +106,10 @@ struct NVPTXLowerArgsPass : PassInfoMixin<NVPTXLowerArgsPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
+struct NVPTXTagInvariantLoadsPass : PassInfoMixin<NVPTXTagInvariantLoadsPass> {
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 namespace NVPTX {
 enum DrvInterface {
   NVCL,

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -766,46 +766,12 @@ NVPTX::Scope NVPTXDAGToDAGISel::getOperationScope(MemSDNode *N,
   llvm_unreachable("unhandled ordering");
 }
 
-static bool canLowerToLDG(MemSDNode *N, const NVPTXSubtarget &Subtarget,
-                          unsigned CodeAddrSpace, MachineFunction *F) {
+static bool canLowerToLDG(const MemSDNode *N, const NVPTXSubtarget &Subtarget,
+                          unsigned CodeAddrSpace) {
   // We use ldg (i.e. ld.global.nc) for invariant loads from the global address
   // space.
-  //
-  // We have two ways of identifying invariant loads: Loads may be explicitly
-  // marked as invariant, or we may infer them to be invariant.
-  //
-  // We currently infer invariance for loads from
-  //  - constant global variables, and
-  //  - kernel function pointer params that are noalias (i.e. __restrict) and
-  //    never written to.
-  //
-  // TODO: Perform a more powerful invariance analysis (ideally IPO, and ideally
-  // not during the SelectionDAG phase).
-  //
-  // TODO: Infer invariance only at -O2.  We still want to use ldg at -O0 for
-  // explicitly invariant loads because these are how clang tells us to use ldg
-  // when the user uses a builtin.
-  if (!Subtarget.hasLDG() || CodeAddrSpace != NVPTX::AddressSpace::Global)
-    return false;
-
-  if (N->isInvariant())
-    return true;
-
-  bool IsKernelFn = isKernelFunction(F->getFunction());
-
-  // We use getUnderlyingObjects() here instead of getUnderlyingObject() mainly
-  // because the former looks through phi nodes while the latter does not. We
-  // need to look through phi nodes to handle pointer induction variables.
-  SmallVector<const Value *, 8> Objs;
-  getUnderlyingObjects(N->getMemOperand()->getValue(), Objs);
-
-  return all_of(Objs, [&](const Value *V) {
-    if (auto *A = dyn_cast<const Argument>(V))
-      return IsKernelFn && A->onlyReadsMemory() && A->hasNoAliasAttr();
-    if (auto *GV = dyn_cast<const GlobalVariable>(V))
-      return GV->isConstant();
-    return false;
-  });
+  return Subtarget.hasLDG() && CodeAddrSpace == NVPTX::AddressSpace::Global &&
+         N->isInvariant();
 }
 
 static unsigned int getFenceOp(NVPTX::Ordering O, NVPTX::Scope S,
@@ -1106,10 +1072,9 @@ bool NVPTXDAGToDAGISel::tryLoad(SDNode *N) {
     return false;
 
   // Address Space Setting
-  unsigned int CodeAddrSpace = getCodeAddrSpace(LD);
-  if (canLowerToLDG(LD, *Subtarget, CodeAddrSpace, MF)) {
+  const unsigned CodeAddrSpace = getCodeAddrSpace(LD);
+  if (canLowerToLDG(LD, *Subtarget, CodeAddrSpace))
     return tryLDGLDU(N);
-  }
 
   SDLoc DL(N);
   SDValue Chain = N->getOperand(0);
@@ -1192,10 +1157,9 @@ bool NVPTXDAGToDAGISel::tryLoadVector(SDNode *N) {
   const MVT MemVT = MemEVT.getSimpleVT();
 
   // Address Space Setting
-  unsigned int CodeAddrSpace = getCodeAddrSpace(MemSD);
-  if (canLowerToLDG(MemSD, *Subtarget, CodeAddrSpace, MF)) {
+  const unsigned CodeAddrSpace = getCodeAddrSpace(MemSD);
+  if (canLowerToLDG(MemSD, *Subtarget, CodeAddrSpace))
     return tryLDGLDU(N);
-  }
 
   EVT EltVT = N->getValueType(0);
   SDLoc DL(N);

llvm/lib/Target/NVPTX/NVPTXPassRegistry.def

@@ -38,5 +38,6 @@ FUNCTION_ALIAS_ANALYSIS("nvptx-aa", NVPTXAA())
 #endif
 FUNCTION_PASS("nvvm-intr-range", NVVMIntrRangePass())
 FUNCTION_PASS("nvptx-copy-byval-args", NVPTXCopyByValArgsPass())
-FUNCTION_PASS("nvptx-lower-args", NVPTXLowerArgsPass(*this));
+FUNCTION_PASS("nvptx-lower-args", NVPTXLowerArgsPass(*this))
+FUNCTION_PASS("nvptx-tag-invariant-loads", NVPTXTagInvariantLoadsPass())
 #undef FUNCTION_PASS

llvm/lib/Target/NVPTX/NVPTXTagInvariantLoads.cpp

@@ -0,0 +1,104 @@
+//===------ NVPTXTagInvariantLoads.cpp - Tag invariant loads --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements invaraint load tagging. It traverses load instructions
+// in a function, and determines if each load can be tagged as invariant.
+//
+// We currently infer invariance for loads from
+//  - constant global variables, and
+//  - kernel function pointer params that are noalias (i.e. __restrict) and
+//    never written to.
+//
+// TODO: Perform a more powerful invariance analysis (ideally IPO).
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTXUtilities.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/Support/NVPTXAddrSpace.h"
+
+using namespace llvm;
+
+static void markLoadsAsInvariant(LoadInst *LI) {
+  LI->setMetadata(LLVMContext::MD_invariant_load,
+                  MDNode::get(LI->getContext(), {}));
+}
+
+static bool tagInvariantLoads(Function &F) {
+  const bool IsKernelFn = isKernelFunction(F);
+
+  bool Changed = false;
+  for (auto &I : instructions(F)) {
+    if (auto *LI = dyn_cast<LoadInst>(&I)) {
+
+      // Don't bother with non-global loads
+      if (LI->getPointerAddressSpace() != NVPTXAS::ADDRESS_SPACE_GLOBAL)
+        continue;
+
+      if (LI->getMetadata(LLVMContext::MD_invariant_load))
+        continue;
+
+      SmallVector<const Value *, 8> Objs;
+
+      // We use getUnderlyingObjects() here instead of getUnderlyingObject()
+      // mainly because the former looks through phi nodes while the latter does
+      // not. We need to look through phi nodes to handle pointer induction
+      // variables.
+
+      getUnderlyingObjects(LI->getPointerOperand(), Objs);
+
+      const bool IsInvariant = all_of(Objs, [&](const Value *V) {
+        if (const auto *A = dyn_cast<const Argument>(V))
+          return IsKernelFn && A->onlyReadsMemory() && A->hasNoAliasAttr();
+        if (const auto *GV = dyn_cast<const GlobalVariable>(V))
+          return GV->isConstant();
+        return false;
+      });
+
+      if (IsInvariant) {
+        markLoadsAsInvariant(LI);
+        Changed = true;
+      }
+    }
+  }
+
+  return Changed;
+}
+
+namespace {
+
+struct NVPTXTagInvariantLoadLegacyPass : public FunctionPass {
+  static char ID;
+
+  NVPTXTagInvariantLoadLegacyPass() : FunctionPass(ID) {}
+  bool runOnFunction(Function &F) override;
+};
+
+} // namespace
+
+INITIALIZE_PASS(NVPTXTagInvariantLoadLegacyPass, "nvptx-tag-invariant-loads",
+                "NVPTX Tag Invariant Loads", false, false)
+
+bool NVPTXTagInvariantLoadLegacyPass::runOnFunction(Function &F) {
+  return tagInvariantLoads(F);
+}
+
+char NVPTXTagInvariantLoadLegacyPass::ID = 0;
+
+FunctionPass *llvm::createNVPTXTagInvariantLoadsPass() {
+  return new NVPTXTagInvariantLoadLegacyPass();
+}
+
+PreservedAnalyses NVPTXTagInvariantLoadsPass::run(Function &F,
+                                                  FunctionAnalysisManager &) {
+  return tagInvariantLoads(F) ? PreservedAnalyses::none()
+                              : PreservedAnalyses::all();
+}

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -112,6 +112,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   initializeNVPTXAAWrapperPassPass(PR);
   initializeNVPTXExternalAAWrapperPass(PR);
   initializeNVPTXPeepholePass(PR);
+  initializeNVPTXTagInvariantLoadLegacyPassPass(PR);
 }
 
 static std::string computeDataLayout(bool is64Bit, bool UseShortPointers) {
@@ -395,6 +396,7 @@ void NVPTXPassConfig::addIRPasses() {
     if (!DisableLoadStoreVectorizer)
       addPass(createLoadStoreVectorizerPass());
     addPass(createSROAPass());
+    addPass(createNVPTXTagInvariantLoadsPass());
   }
 
   if (ST.hasPTXASUnreachableBug()) {