[AMDGPU][AMDGPULowerKernelAttributes] Use stripAndAccumulateConstantOffset

[jmmartinez]

Summary

This removes the need to handle explicitely the presence or absence of bitcasts when typed pointers are used. It also takes into account the case where we have more than one use of load(call @llvm.amdgc.implicitarg.ptr())

[jmmartinez]

• Handle non-inbounds GEP

[jdoerfert]

Some drive by comments

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

84

This function and uses are unrelated.

89

This does break with the "unique" property below. Two users, one simple one not, will result in a load. Is that expected?

105

A unique load reached via select or phi will not be detected. UniqueLoad && UniqueLoad != L can help with that.

Style: I'd remove braces for single return conditionals.

138

I would use the Attributor::checkForAllUses, also unsure why the loads need to be unique. Depending on the call AAPointerInfo would allow you to handle way more things without any of this manual tracing. Either seems to be a useful rewrite beyond the scope of this. Feel free to ignore for now.

[jmmartinez]

• Rename getUniqueLoadUser into getUniqueSimpleLoadUser.
• Consider the case where the UniqueLoad is 'found' twice.
• Style

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

84

I'm sorry, I do not understand this remark. What do you mean?

89

It's ok if there is another non-simple load.

In fact, this pass could be modified to support multiple loads. My guess is that for simplicity only one is supported for the moment.

I renamed the function instead into getUniqueSimpleLoadUser

138

I'm not familiar with it. I'm going to take a look at it. Thanks!

[jdoerfert]

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

84

The commit is about the load tracking and stripping. This change (outlining getBasePtrIntrinsic and using it) is unrelated to the content/description of the patch. Make it two patches instead.

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

84

Oh! I agree. I created https://reviews.llvm.org/D135266 for that change.

[arsenm]

Needs a test for this multiple load case

[jmmartinez]

• Rebase and move annonymous namespace to remove static keyword

[jmmartinez]

• Extend AMDGPULowerKernelAttributes to handle multiple occurrences of the attributes in the same function
• Fix reminder indices in test
• Add test-cases with multiple users, and more than one in a single function

[jmmartinez]

To split into smaller changes later.

llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp

public:
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
  }
 };
 
 } // end anonymous namespace
 
+static Function *getBasePtrIntrinsic(Module &M, bool IsV5OrAbove) {
+  auto Intrinsic = IsV5OrAbove ? Intrinsic::amdgcn_implicitarg_ptr
+                               : Intrinsic::amdgcn_dispatch_ptr;
+  StringRef Name = Intrinsic::getName(Intrinsic);
+  return M.getFunction(Name);
+}
+
+static LoadInst *getUniqueLoadUser(Value *V) {
+  if (LoadInst *L = dyn_cast<LoadInst>(V)) {
+    return L->isSimple() ? L : nullptr;
+  }
+
+  if (!V->getType()->isPointerTy()) {
+    return nullptr;
+  }
+
+  LoadInst *UniqueLoad = nullptr;
+  for (User *U : V->users()) {
+    if (LoadInst *L = getUniqueLoadUser(U)) {
+      if (UniqueLoad) {
+        return nullptr;
+      }
+      UniqueLoad = L;
+    }
+  }
+  return UniqueLoad;
+}
+
 static bool processUse(CallInst *CI, bool IsV5OrAbove) {
-  Function *F = CI->getParent()->getParent();
+  Function *F = CI->getFunction();
 
   auto MD = F->getMetadata("reqd_work_group_size");
   const bool HasReqdWorkGroupSize = MD && MD->getNumOperands() == 3;
 
   const bool HasUniformWorkGroupSize =
     F->getFnAttribute("uniform-work-group-size").getValueAsBool();
 
   if (!HasReqdWorkGroupSize && !HasUniformWorkGroupSize)
     return false;
 
   Value *BlockCounts[3] = {nullptr, nullptr, nullptr};
   Value *GroupSizes[3]  = {nullptr, nullptr, nullptr};
   Value *Remainders[3]  = {nullptr, nullptr, nullptr};
   Value *GridSizes[3]   = {nullptr, nullptr, nullptr};
 
   const DataLayout &DL = F->getParent()->getDataLayout();
 
   // We expect to see several GEP users, casted to the appropriate type and
   // loaded.
   for (User *U : CI->users()) {
-    if (!U->hasOneUse())
-      continue;
-
-    int64_t Offset = 0;
-    auto *Load = dyn_cast<LoadInst>(U); // Load from ImplicitArgPtr/DispatchPtr?
-    auto *BCI = dyn_cast<BitCastInst>(U);
-    if (!Load && !BCI) {
-      if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
-        continue;
-      Load = dyn_cast<LoadInst>(*U->user_begin()); // Load from GEP?
-      BCI = dyn_cast<BitCastInst>(*U->user_begin());
+    LoadInst *Load = getUniqueLoadUser(U);
+    if (!Load) {
     }
 
-    if (BCI) {
-      if (!BCI->hasOneUse())
-        continue;
-      Load = dyn_cast<LoadInst>(*BCI->user_begin()); // Load from BCI?
+    APInt Offset(64, 0U);
+    if (Load != U &&
+        U->stripAndAccumulateConstantOffsets(DL, Offset, true) != CI) {
     }
 
-    if (!Load || !Load->isSimple())
-      continue;
-
     unsigned LoadSize = DL.getTypeStoreSize(Load->getType());
 
     // TODO: Handle merged loads.
+    auto const OffsetValue = Offset.getSExtValue();
     if (IsV5OrAbove) { // Base is ImplicitArgPtr.
-      switch (Offset) {
+      switch (OffsetValue) {
       case HIDDEN_BLOCK_COUNT_X:
         if (LoadSize == 4)
           BlockCounts[0] = Load;
         break;
       case HIDDEN_BLOCK_COUNT_Y:
         if (LoadSize == 4)
           BlockCounts[1] = Load;
         break;
if (IsV5OrAbove) { // Base is ImplicitArgPtr.
       case HIDDEN_REMAINDER_Z:
         if (LoadSize == 2)
           Remainders[2] = Load;
         break;
       default:
         break;
       }
     } else { // Base is DispatchPtr.
-      switch (Offset) {
+      switch (OffsetValue) {
       case WORKGROUP_SIZE_X:
         if (LoadSize == 2)
           GroupSizes[0] = Load;
         break;
       case WORKGROUP_SIZE_Y:
         if (LoadSize == 2)
           GroupSizes[1] = Load;
         break;
static bool processUse(CallInst *CI, bool IsV5OrAbove) {
   return MadeChange;
 }
 
 
 // TODO: Move makeLIDRangeMetadata usage into here. Seem to not get
 // TargetPassConfig for subtarget.
 bool AMDGPULowerKernelAttributes::runOnModule(Module &M) {
   bool MadeChange = false;
-  Function *BasePtr = nullptr;
   bool IsV5OrAbove = AMDGPU::getAmdhsaCodeObjectVersion() >= 5;
-  if (IsV5OrAbove) {
-    StringRef ImplicitArgPtrName =
-        Intrinsic::getName(Intrinsic::amdgcn_implicitarg_ptr);
-    BasePtr = M.getFunction(ImplicitArgPtrName);
-  } else { // Pre-V5.
-    StringRef DispatchPtrName =
-        Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);
-    BasePtr = M.getFunction(DispatchPtrName);
-  }
+  Function *BasePtr = getBasePtrIntrinsic(M, IsV5OrAbove);
 
   if (!BasePtr) // ImplicitArgPtr/DispatchPtr not used.
     return false;
 
   SmallPtrSet<Instruction *, 4> HandledUses;
   for (auto *U : BasePtr->users()) {
     CallInst *CI = cast<CallInst>(U);
     if (HandledUses.insert(CI).second) {

 char AMDGPULowerKernelAttributes::ID = 0;
 
 ModulePass *llvm::createAMDGPULowerKernelAttributesPass() {
   return new AMDGPULowerKernelAttributes();
 }
 
 PreservedAnalyses
 AMDGPULowerKernelAttributesPass::run(Function &F, FunctionAnalysisManager &AM) {
-  Function *BasePtr = nullptr;
   bool IsV5OrAbove = AMDGPU::getAmdhsaCodeObjectVersion() >= 5;
-  if (IsV5OrAbove) {
-    StringRef ImplicitArgPtrName =
-        Intrinsic::getName(Intrinsic::amdgcn_implicitarg_ptr);
-    BasePtr = F.getParent()->getFunction(ImplicitArgPtrName);
-  } else { // Pre_V5.
-    StringRef DispatchPtrName =
-        Intrinsic::getName(Intrinsic::amdgcn_dispatch_ptr);
-    BasePtr = F.getParent()->getFunction(DispatchPtrName);
-  }
+  Function *BasePtr = getBasePtrIntrinsic(*F.getParent(), IsV5OrAbove);
 
   if (!BasePtr) // ImplicitArgPtr/DispatchPtr not used.
     return PreservedAnalyses::all();
 
   for (Instruction &I : instructions(F)) {
     if (CallInst *CI = dyn_cast<CallInst>(&I)) {
       if (CI->getCalledFunction() == BasePtr)
         processUse(CI, IsV5OrAbove);
     }
   }
 
   return PreservedAnalyses::all();
 }