Preserve inbounds information of GEP during Argument Promotion Pass across callee and its callers.

[pravinjagtap]

Summary

Presently inbounds information is lost between callee and its callers when Argument is promoted.
This change respects the inbounds of GEP emitted by front-end for callee and emits GEP with preserved
inbounds information in callers instead taking conservative approach of emitting GEP without inbounds
by default.

Note that Argument is only promoted when its *only* loaded (i.e. no store) in Argument Promotion Pass.

Motivation for this change is that the AMDGPUPromoteKernelArgumentsPass (https://reviews.llvm.org/D111464)
can infer the address-spaces and promote FLAT memory operation to GLOBAL based on this change.

[nikic]

It looks like you need to update a number of existing tests.

You should also test the case where the GEP inbounds is conditional -- I think it might still be okay based on a dereferenceabiliy argument, but it's not entirely obvious to me.

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

283

I think using inbounds here is okay, but the justification is not correct. You probably want to argue about byval-implied dereferenceability here.

[pravinjagtap]

In D119127#3300894, @nikic wrote:

It looks like you need to update a number of existing tests.

You should also test the case where the GEP inbounds is conditional -- I think it might still be okay based on a dereferenceabiliy argument, but it's not entirely obvious to me.

Thank you for your inputs. I am not sure whether I fully understood what do you mean by GEP inbounds is conditional. I am assuming that test should include: GEP with inbound and without inbound in callee so that other conditional branch will be exercised.

[pravinjagtap]

Addressed the comments by nikic and fixed the unit tests

[nikic]

In D119127#3304281, @pravinjagtap wrote:

In D119127#3300894, @nikic wrote:

It looks like you need to update a number of existing tests.

You should also test the case where the GEP inbounds is conditional -- I think it might still be okay based on a dereferenceabiliy argument, but it's not entirely obvious to me.

Thank you for your inputs. I am not sure whether I fully understood what do you mean by GEP inbounds is conditional. I am assuming that test should include: GEP with inbound and without inbound in callee so that other conditional branch will be exercised.

Sorry for the confusion: What I meant is a case where the GEP + load is not guaranteed to execute on entry to the function, but is behind a conditional branch.

[pravinjagtap]

Applied git-clang-format on change. Sorry for inconvenience.

[pravinjagtap]

Added test which exercises the conditional GEP inbound

[pravinjagtap]

In D119127#3304281, @pravinjagtap wrote:

In D119127#3300894, @nikic wrote:

It looks like you need to update a number of existing tests.

You should also test the case where the GEP inbounds is conditional -- I think it might still be okay based on a dereferenceabiliy argument, but it's not entirely obvious to me.

Thank you for your inputs. I am not sure whether I fully understood what do you mean by GEP inbounds is conditional. I am assuming that test should include: GEP with inbound and without inbound in callee so that other conditional branch will be exercised.

Added test /llvm/test/Transforms/ArgumentPromotion/conditional-gep-and-load.ll which exercises the conditional inbounds @nikic.

[vpykhtin]

Using IsArgGepInBound map just to hold a boolean seems expensive to me, can you just add a field to OriginalLoads value?

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

201

UI isn't expected to be null here, dyn_cast should be enough.

[pravinjagtap]

Addressed vpykhtin's commnets

[vpykhtin]

I would replace pair of pair with ordinary struct, much easier to read the code

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

doPromotion(Function *F, SmallPtrSetImpl<Argument *> &ArgsToPromote,
 
   // OriginalLoads - Keep track of a representative load instruction from the
   // original function so that we can tell the alias analysis implementation
   // what the new GEP/Load instructions we are inserting look like.
   // We need to keep the original loads for each argument and the elements
   // of the argument that are accessed.
   std::map<std::pair<Argument *, IndicesVector>, LoadInst *> OriginalLoads;
 
+  // Keeps track of inbounds information of GEP instruction used to load
+  // promoted arguments present in callee function. This is used to preserve
+  // this inbound information in callers as well.
+  std::map<std::pair<Argument *, IndicesVector>, bool> IsArgGepInBound;
+
   // Attribute - Keep track of the parameter attributes for the arguments
   // that we are *not* promoting. For the ones that we do promote, the parameter
   // attributes are lost
   SmallVector<AttributeSet, 8> ArgAttrVec;
   AttributeList PAL = F->getAttributes();
 
   // First, determine the new argument list
   unsigned ArgNo = 0;
if (ByValArgsToTransform.count(&*I)) {
         ArgIndices.insert(std::make_pair(SrcTy, Indices));
         LoadInst *OrigLoad;
         if (LoadInst *L = dyn_cast<LoadInst>(UI))
           OrigLoad = L;
         else
           // Take any load, we will use it only to update Alias Analysis
           OrigLoad = cast<LoadInst>(UI->user_back());
         OriginalLoads[std::make_pair(&*I, Indices)] = OrigLoad;
+        // Preserve inbounds information of GEP so that it can be used in
+        // callers of current function
+        const GEPOperator *GEPOp = dyn_cast_or_null<GEPOperator>(UI);
+        if (GEPOp) {
+          // Indices are unique in GEP since all the dead GEPs are already
+          // removed
+          IsArgGepInBound[std::make_pair(&*I, Indices)] = GEPOp->isInBounds();
+        }
       }
 
       // Add a parameter to the function for each element passed in.
       for (const auto &ArgIndex : ArgIndices) {
         // not allowed to dereference ->begin() if size() is 0
         Params.push_back(GetElementPtrInst::getIndexedType(
             I->getType()->getPointerElementType(), ArgIndex.second));
         ArgAttrVec.push_back(AttributeSet());
for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
         Type *AgTy = I->getParamByValType();
         StructType *STy = cast<StructType>(AgTy);
         Value *Idxs[2] = {
             ConstantInt::get(Type::getInt32Ty(F->getContext()), 0), nullptr};
         const StructLayout *SL = DL.getStructLayout(STy);
         Align StructAlign = *I->getParamAlign();
         for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i) {
           Idxs[1] = ConstantInt::get(Type::getInt32Ty(F->getContext()), i);
-          auto *Idx =
-              IRB.CreateGEP(STy, *AI, Idxs, (*AI)->getName() + "." + Twine(i));
+          // It is safe to create inbounds GEP since source type is struct
+          auto *Idx = IRB.CreateInBoundsGEP(STy, *AI, Idxs,
+                                            (*AI)->getName() + "." + Twine(i));
           // TODO: Tell AA about the new values?
           Align Alignment =
               commonAlignment(StructAlign, SL->getElementOffset(i));
           Args.push_back(IRB.CreateAlignedLoad(
               STy->getElementType(i), Idx, Alignment, Idx->getName() + ".val"));
           ArgAttrVec.push_back(AttributeSet());
         }
       } else if (!I->use_empty()) {
for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
               Ops.push_back(ConstantInt::get(IdxTy, II));
               // Keep track of the type we're currently indexing.
               if (auto *ElPTy = dyn_cast<PointerType>(ElTy))
                 ElTy = ElPTy->getPointerElementType();
               else
                 ElTy = GetElementPtrInst::getTypeAtIndex(ElTy, II);
             }
             // And create a GEP to extract those indices.
+            bool IsCalleeArgGepInBound =
+                IsArgGepInBound[std::make_pair(&*I, ArgIndex.second)];
+            if (IsCalleeArgGepInBound) {
+              // Using preserved inbounds information of GEP in callee while
+              // creating GEP in caller.
+              V = IRB.CreateInBoundsGEP(ArgIndex.first, V, Ops,
+                                        V->getName() + ".idx");
+            } else {
               V = IRB.CreateGEP(ArgIndex.first, V, Ops, V->getName() + ".idx");
+            }
             Ops.clear();
           }
           // Since we're replacing a load make sure we take the alignment
           // of the previous load.
           LoadInst *newLoad =
               IRB.CreateLoad(OrigLoad->getType(), V, V->getName() + ".val");
           newLoad->setAlignment(OrigLoad->getAlign());
           // Transfer the AA info too.

llvm/test/Transforms/ArgumentPromotion/preserve-inbounds-info-of-gep.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --scrub-attributes
+; RUN: opt < %s -passes=argpromotion -S | FileCheck %s
+
+%T = type { i32, i32, i32, i32 }
+@G = constant %T { i32 0, i32 0, i32 17, i32 25 }
+
+define internal i32 @test(%T* %p) {
+; CHECK-LABEL: define {{[^@]+}}@test
+; CHECK-SAME: (i32 [[P_0_2_VAL:%.*]], i32 [[P_0_3_VAL:%.*]])
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = add i32 [[P_0_3_VAL]], [[P_0_2_VAL]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %a.gep = getelementptr inbounds %T, %T* %p, i64 0, i32 3
+  %b.gep = getelementptr inbounds %T, %T* %p, i64 0, i32 2
+  %a = load i32, i32* %a.gep
+  %b = load i32, i32* %b.gep
+  %v = add i32 %a, %b
+  ret i32 %v
+}
+
+define i32 @caller() {
+; CHECK-LABEL: define {{[^@]+}}@caller()
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[G_IDX:%.*]] = getelementptr inbounds [[T:%.*]], %T* @G, i64 0, i32 2
+; CHECK-NEXT:    [[G_IDX_VAL:%.*]] = load i32, i32* [[G_IDX]]
+; CHECK-NEXT:    [[G_IDX1:%.*]] = getelementptr inbounds [[T]], %T* @G, i64 0, i32 3
+; CHECK-NEXT:    [[G_IDX1_VAL:%.*]] = load i32, i32* [[G_IDX1]]
+; CHECK-NEXT:    [[V:%.*]] = call i32 @test(i32 [[G_IDX_VAL]], i32 [[G_IDX1_VAL]])
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %v = call i32 @test(%T* @G)
+  ret i32 %v
+}