[ArgPromotion] Update allocsize indices after promotion

Promotion can add/remove arguments. We need to update the
indices in the allocsize attribute accordingly.

Fixes https://github.com/llvm/llvm-project/issues/66103.

1 files changed, 21 insertions, 1 deletions

diff --git a/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp b/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
index 824da63..db9d5bc 100644
--- a/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
+++ b/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
@@ -121,19 +121,24 @@ doPromotion(Function *F, FunctionAnalysisManager &FAM,
   // that we are *not* promoting. For the ones that we do promote, the parameter
   // attributes are lost
   SmallVector<AttributeSet, 8> ArgAttrVec;
+  // Mapping from old to new argument indices. -1 for promoted or removed
+  // arguments.
+  SmallVector<unsigned> NewArgIndices;
   AttributeList PAL = F->getAttributes();
 
   // First, determine the new argument list
-  unsigned ArgNo = 0;
+  unsigned ArgNo = 0, NewArgNo = 0;
   for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end(); I != E;
        ++I, ++ArgNo) {
     if (!ArgsToPromote.count(&*I)) {
       // Unchanged argument
       Params.push_back(I->getType());
       ArgAttrVec.push_back(PAL.getParamAttrs(ArgNo));
+      NewArgIndices.push_back(NewArgNo++);
     } else if (I->use_empty()) {
       // Dead argument (which are always marked as promotable)
       ++NumArgumentsDead;
+      NewArgIndices.push_back((unsigned)-1);
     } else {
       const auto &ArgParts = ArgsToPromote.find(&*I)->second;
       for (const auto &Pair : ArgParts) {
@@ -141,6 +146,8 @@ doPromotion(Function *F, FunctionAnalysisManager &FAM,
         ArgAttrVec.push_back(AttributeSet());
       }
       ++NumArgumentsPromoted;
+      NewArgIndices.push_back((unsigned)-1);
+      NewArgNo += ArgParts.size();
     }
   }
 
@@ -173,6 +180,19 @@ doPromotion(Function *F, FunctionAnalysisManager &FAM,
   // the function.
   NF->setAttributes(AttributeList::get(F->getContext(), PAL.getFnAttrs(),
                                        PAL.getRetAttrs(), ArgAttrVec));
+
+  // Remap argument indices in allocsize attribute.
+  if (auto AllocSize = NF->getAttributes().getFnAttrs().getAllocSizeArgs()) {
+    unsigned Arg1 = NewArgIndices[AllocSize->first];
+    assert(Arg1 != (unsigned)-1 && "allocsize cannot be promoted argument");
+    std::optional<unsigned> Arg2;
+    if (AllocSize->second) {
+      Arg2 = NewArgIndices[*AllocSize->second];
+      assert(Arg2 != (unsigned)-1 && "allocsize cannot be promoted argument");
+    }
+    NF->addFnAttr(Attribute::getWithAllocSizeArgs(F->getContext(), Arg1, Arg2));
+  }
+
   AttributeFuncs::updateMinLegalVectorWidthAttr(*NF, LargestVectorWidth);
   ArgAttrVec.clear();