address comments

Author: AlexMaclean

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -766,12 +766,12 @@ NVPTX::Scope NVPTXDAGToDAGISel::getOperationScope(MemSDNode *N,
   llvm_unreachable("unhandled ordering");
 }
 
-static bool canLowerToLDG(const MemSDNode *N, const NVPTXSubtarget &Subtarget,
+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.
   return Subtarget.hasLDG() && CodeAddrSpace == NVPTX::AddressSpace::Global &&
-         N->isInvariant();
+         N.isInvariant();
 }
 
 static unsigned int getFenceOp(NVPTX::Ordering O, NVPTX::Scope S,
@@ -1073,7 +1073,7 @@ bool NVPTXDAGToDAGISel::tryLoad(SDNode *N) {
 
   // Address Space Setting
   const unsigned CodeAddrSpace = getCodeAddrSpace(LD);
-  if (canLowerToLDG(LD, *Subtarget, CodeAddrSpace))
+  if (canLowerToLDG(*LD, *Subtarget, CodeAddrSpace))
     return tryLDGLDU(N);
 
   SDLoc DL(N);
@@ -1158,7 +1158,7 @@ bool NVPTXDAGToDAGISel::tryLoadVector(SDNode *N) {
 
   // Address Space Setting
   const unsigned CodeAddrSpace = getCodeAddrSpace(MemSD);
-  if (canLowerToLDG(MemSD, *Subtarget, CodeAddrSpace))
+  if (canLowerToLDG(*MemSD, *Subtarget, CodeAddrSpace))
     return tryLDGLDU(N);
 
   EVT EltVT = N->getValueType(0);

llvm/lib/Target/NVPTX/NVPTXTagInvariantLoads.cpp

@@ -27,6 +27,32 @@
 
 using namespace llvm;
 
+static bool isInvariantLoad(const LoadInst *LI, const bool IsKernelFn) {
+  // Don't bother with non-global loads
+  if (LI->getPointerAddressSpace() != NVPTXAS::ADDRESS_SPACE_GLOBAL)
+    return false;
+
+  // If the load is already marked as invariant, we don't need to do anything
+  if (LI->getMetadata(LLVMContext::MD_invariant_load))
+    return false;
+
+  // 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(LI->getPointerOperand(), Objs);
+
+  return all_of(Objs, [&](const Value *V) {
+    if (const auto *A = dyn_cast<const Argument>(V))
+      return IsKernelFn && ((A->onlyReadsMemory() && A->hasNoAliasAttr()) ||
+                            isParamGridConstant(*A));
+    if (const auto *GV = dyn_cast<const GlobalVariable>(V))
+      return GV->isConstant();
+    return false;
+  });
+}
+
 static void markLoadsAsInvariant(LoadInst *LI) {
   LI->setMetadata(LLVMContext::MD_invariant_load,
                   MDNode::get(LI->getContext(), {}));
@@ -38,38 +64,12 @@ static bool tagInvariantLoads(Function &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) {
+      if (isInvariantLoad(LI, IsKernelFn)) {
         markLoadsAsInvariant(LI);
         Changed = true;
       }
     }
   }
-
   return Changed;
 }