[ArgPromotion] Replace all md uses of promoted values with undef.

[fhahn]

Summary

Currently, ArgPromotion may leave metadata uses of promoted values,
which will end up in the wrong function, creating invalid IR.

PR33641 fixed this for dead arguments, but it can be also be triggered
arguments with users that are promoted (see the updated test case).

We also have to drop uses to them after promoting them. We need to do
this after dealing with the non-metadata uses, so I also moved the empty
use case to the loop that deals with updating the arguments of the new
function.

[aprantl]

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

419

We could use

llvm::scope_exit RauwUndef([&]() {
// There potentially are metadata uses left for things like
// llvm.dbg.value. Replace them with undef.
I->replaceAllUsesWith(UndefValue::get(I->getType()));
});

if (I->use_empty())
  continue;

I'll leave it up to you to decide whether that's actually better.

[aprantl]

Thanks!

[fhahn]

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

419

Nice, I wasn't aware of scope_exit()! That would allow to get rid of the unfortunate duplication.

The reason we need to drop at the end of the scope is that there is no replaceAllMetadataUsesWith and metadata uses can only be replaced together with regular uses. So alternatively we could also add a replaceAllMetadataUsesWith, but if that's the only place it's probably not worth it. What do you think?

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
if (ByValArgsToTransform.count(&*I)) {
       ++NumByValArgsPromoted;
     } else if (!ArgsToPromote.count(&*I)) {
       // Unchanged argument
       Params.push_back(I->getType());
       ArgAttrVec.push_back(PAL.getParamAttributes(ArgNo));
     } else if (I->use_empty()) {
       // Dead argument (which are always marked as promotable)
       ++NumArgumentsDead;
-
-      // There may be remaining metadata uses of the argument for things like
-      // llvm.dbg.value. Replace them with undef.
-      I->replaceAllUsesWith(UndefValue::get(I->getType()));
     } else {
       // Okay, this is being promoted. This means that the only uses are loads
       // or GEPs which are only used by loads
 
       // In this table, we will track which indices are loaded from the argument
       // (where direct loads are tracked as no indices).
       ScalarizeTable &ArgIndices = ScalarizedElements[&*I];
       for (User *U : I->users()) {
if (ByValArgsToTransform.count(&*I)) {
         CallInst *Call = dyn_cast<CallInst>(U);
         if (!Call)
           continue;
         Call->setTailCall(false);
       }
       continue;
     }
 
+    // There potentially are metadata uses for things like llvm.dbg.value.
+    // Replace them with undef, after handling the other regular uses.
+    auto RauwUndefMetadata = make_scope_exit(
+        [&]() { I->replaceAllUsesWith(UndefValue::get(I->getType())); });
+
     if (I->use_empty())
       continue;
 
     // Otherwise, if we promoted this argument, then all users are load
     // instructions (or GEPs with only load users), and all loads should be
     // using the new argument that we added.
     ScalarizeTable &ArgIndices = ScalarizedElements[&*I];
 
     while (!I->use_empty()) {
while (!I->use_empty()) {
         while (!GEP->use_empty()) {
           LoadInst *L = cast<LoadInst>(GEP->user_back());
           L->replaceAllUsesWith(&*TheArg);
           L->eraseFromParent();
         }
         GEP->eraseFromParent();
       }
     }
-
     // Increment I2 past all of the arguments added for this promoted pointer.
     std::advance(I2, ArgIndices.size());
   }
 
   return NF;
 }
 
 /// Return true if we can prove that all callees pass in a valid pointer for the

llvm/test/Transforms/ArgumentPromotion/pr33641_remove_arg_dbgvalue.ll

 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.dbg.value(metadata %p_t %p, metadata !4, metadata !5), !dbg !6
   ret void
 }
 
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
+
+; Test case where the promoted argument has uses in @callee and we need to
+; retain a reference to the original function, because it is stored in @storer.
+define void @storer({i32, i32}* %ptr) {
+; CHECK-LABEL: define {{[^@]+}}@storer
+; CHECK-SAME: ({ i32, i32 }* [[PTR:%.*]])
+; CHECK-NEXT:    ret void
+;
+  %tmp = alloca i32 ({i32, i32}*)*
+  store i32 ({i32, i32}*)* @callee,  i32 ({i32, i32}*)** %tmp
+  ret void
+}
+
+define i32 @caller() {
+; CHECK-LABEL: define {{[^@]+}}@caller()
+; CHECK-NEXT:    [[TMP:%.*]] = alloca { i32, i32 }, align 8
+; CHECK-NEXT:    [[F_1:%.*]] = getelementptr { i32, i32 }, { i32, i32 }* [[TMP]], i32 0, i32 1
+; CHECK-NEXT:    store i32 10, i32* [[F_1]], align 4
+; CHECK-NEXT:    [[TMP_IDX:%.*]] = getelementptr { i32, i32 }, { i32, i32 }* [[TMP]], i64 0, i32 1
+; CHECK-NEXT:    [[TMP_IDX_VAL:%.*]] = load i32, i32* [[TMP_IDX]], align 4
+; CHECK-NEXT:    [[RES:%.*]] = call i32 @callee(i32 [[TMP_IDX_VAL]])
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %tmp = alloca {i32, i32}
+  %f.1 = getelementptr {i32, i32}, {i32, i32}* %tmp, i32 0, i32 1
+  store i32 10, i32* %f.1
+  %res = call i32 @callee({i32, i32}* %tmp)
+  ret i32 %res
+}
+
+define internal i32 @callee({i32, i32}* %ptr)  !dbg !7 {
+; CHECK-LABEL: define {{[^@]+}}@callee
+; CHECK-SAME: (i32 [[PTR_0_1_VAL:%.*]]) !dbg !6
+; CHECK-NEXT:    call void @llvm.dbg.value(metadata { i32, i32 }* undef, metadata !7, metadata !DIExpression()), !dbg !8
+; CHECK-NEXT:    call void @llvm.dbg.value(metadata i32 [[PTR_0_1_VAL]], metadata !7, metadata !DIExpression()), !dbg !8
+; CHECK-NEXT:    ret i32 [[PTR_0_1_VAL]]
+;
+  call void @llvm.dbg.value(metadata {i32, i32}* %ptr, metadata !8, metadata !9), !dbg !10
+  %f.1 = getelementptr {i32, i32}, {i32, i32}* %ptr, i32 0, i32 1
+  %l.1 = load i32, i32* %f.1
+  call void @llvm.dbg.value(metadata i32 %l.1, metadata !8, metadata !9), !dbg !10
+  ret i32 %l.1
+}
+
+
+
 !llvm.dbg.cu = !{!0}
 !llvm.module.flags = !{!2}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1)
 !1 = !DIFile(filename: "test.c", directory: "")
 !2 = !{i32 2, !"Debug Info Version", i32 3}
 !3 = distinct !DISubprogram(name: "bar", unit: !0)
 !4 = !DILocalVariable(name: "p", scope: !3)
 !5 = !DIExpression()
 !6 = !DILocation(line: 1, column: 1, scope: !3)
+!7 = distinct !DISubprogram(name: "callee", unit: !0)
+!8 = !DILocalVariable(name: "c", scope: !7)
+!9 = !DIExpression()
+!10 = !DILocation(line: 2, column: 2, scope: !7)