[Inliner] Fix noalias metadata handling for instructions simplified during cloning (PR50270)

[nikic]

Summary

Instead of using VMap, which may include instructions from the caller as a result of simplification, iterate over the (FirstNewBlock, Caller->end()) range, which will only include new instructions.

Fixes https://bugs.llvm.org/show_bug.cgi?id=50270.

[nikic]

llvm/lib/Transforms/Utils/InlineFunction.cpp

913

I've left this as a TODO to be more confident that this is safe to backport to LLVM 12.

[jeroen.dobbelaere]

Looks good. Could you also add the other example as a test case ?

I am wondering if we should extend the testcases with the llvm.experimental.noalias.scope.decl. I assume they got lost while reducing the testcases ?

[nikic]

Add scope decls to test.

[nikic]

In D102110#2747213, @jeroen.dobbelaere wrote:

Looks good. Could you also add the other example as a test case ?

Both examples are already included in the test: @caller tests PropagateCallSiteMetadata and @self_caller tests ScopedAliasMetadataDeepCloner::remap.

I am wondering if we should extend the testcases with the llvm.experimental.noalias.scope.decl. I assume they got lost while reducing the testcases ?

Done! I hadn't included them originally because they're not strictly necessary, but it does make sense to have them for clarity.

[jeroen.dobbelaere]

In D102110#2748594, @nikic wrote:

In D102110#2747213, @jeroen.dobbelaere wrote:

Looks good. Could you also add the other example as a test case ?

Both examples are already included in the test: @caller tests PropagateCallSiteMetadata and @self_caller tests ScopedAliasMetadataDeepCloner::remap.

Yes indeed. Not sure how I overlooked it.

I am wondering if we should extend the testcases with the llvm.experimental.noalias.scope.decl. I assume they got lost while reducing the testcases ?

Done! I hadn't included them originally because they're not strictly necessary, but it does make sense to have them for clarity.

Thanks.

LGTM

[jeroen.dobbelaere]

llvm/lib/Transforms/Utils/InlineFunction.cpp

997

Just wondering: don't we have a similar issue here ?

[nikic]

llvm/lib/Transforms/Utils/InlineFunction.cpp

997

Yes, unfortunately we do. It took me a while, but I found a suitable llvm.masked.load fold to exploit:

define <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* noalias %ptr2, <2 x i1> %mask, <2 x i8> %passthru) {
  %ret = call <2 x i8> @llvm.masked.load.v2i8(<2 x i8>* %ptr1, i32 1, <2 x i1> %mask, <2 x i8> %passthru)
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2
  ret <2 x i8> %ret
}

define void @caller(<2 x i8>* %ptr1, <2 x i8>* %ptr2) {
  %passthru = load <2 x i8>, <2 x i8>* %ptr2
  call <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* %ptr2, <2 x i1> zeroinitializer, <2 x i8> %passthru)
  ret void
}

declare <2 x i8> @llvm.masked.load.v2i8(<2 x i8>*, i32, <2 x i1>, <2 x i8>)

Results in:

define <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* noalias %ptr2, <2 x i1> %mask, <2 x i8> %passthru) {
  %ret = call <2 x i8> @llvm.masked.load.v2i8.p0v2i8(<2 x i8>* %ptr1, i32 1, <2 x i1> %mask, <2 x i8> %passthru)
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2, align 2
  ret <2 x i8> %ret
}

define void @caller(<2 x i8>* %ptr1, <2 x i8>* %ptr2) {
  %passthru = load <2 x i8>, <2 x i8>* %ptr2, align 2, !noalias !0
  call void @llvm.experimental.noalias.scope.decl(metadata !0)
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2, align 2, !alias.scope !0
  ret void
}

In this case, we don't even need the incorrect annotation to go on an instruction from the caller, it can also happen with a callee instruction:

define <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* noalias %ptr2, <2 x i1> %mask) {
  %passthru = load <2 x i8>, <2 x i8>* %ptr2
  %ret = call <2 x i8> @llvm.masked.load.v2i8(<2 x i8>* %ptr1, i32 1, <2 x i1> %mask, <2 x i8> %passthru)
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2
  ret <2 x i8> %ret
}

define void @caller(<2 x i8>* %ptr1, <2 x i8>* %ptr2) {
  call <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* %ptr2, <2 x i1> zeroinitializer)
  ret void
}

declare <2 x i8> @llvm.masked.load.v2i8(<2 x i8>*, i32, <2 x i1>, <2 x i8>)

Results in:

define <2 x i8> @callee(<2 x i8>* %ptr1, <2 x i8>* noalias %ptr2, <2 x i1> %mask) {
  %passthru = load <2 x i8>, <2 x i8>* %ptr2, align 2
  %ret = call <2 x i8> @llvm.masked.load.v2i8.p0v2i8(<2 x i8>* %ptr1, i32 1, <2 x i1> %mask, <2 x i8> %passthru)
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2, align 2
  ret <2 x i8> %ret
}

define void @caller(<2 x i8>* %ptr1, <2 x i8>* %ptr2) {
  call void @llvm.experimental.noalias.scope.decl(metadata !0)
  %passthru.i = load <2 x i8>, <2 x i8>* %ptr2, align 2, !alias.scope !0, !noalias !0
  store <2 x i8> zeroinitializer, <2 x i8>* %ptr2, align 2, !alias.scope !0
  ret void
}

Note how the metadata is claiming that the load does not alias ... with itself.

So, in this case the problem is a bit different from the other two: We have to not just prevent annotating instructions from the caller, we should prevent annotating any instructions for which some kind of folding occurred.

[nikic]

llvm/lib/Transforms/Utils/InlineFunction.cpp

997

Anyone know how to fix this nicely? I don't think we have any information on whether an instruction was simplified or not at this point. The only thing that comes to mind is to go through all instructions and check which ones occur in the map multiple times.

Do we need to change the cloning implementation to track simplification info in a side table?

[nikic]

llvm/lib/Transforms/Utils/InlineFunction.cpp

997

I submitted https://bugs.llvm.org/show_bug.cgi?id=50589 to keep track of this issue.

llvm/lib/Transforms/Utils/InlineFunction.cpp

static void HandleInlinedEHPad(InvokeInst *II, BasicBlock *FirstNewBlock,
   // invoke instruction. Eliminate these entries (which might even delete the
   // PHI node) now.
   UnwindDest->removePredecessor(InvokeBB);
 }
 
 /// When inlining a call site that has !llvm.mem.parallel_loop_access,
 /// !llvm.access.group, !alias.scope or !noalias metadata, that metadata should
 /// be propagated to all memory-accessing cloned instructions.
-static void PropagateCallSiteMetadata(CallBase &CB, ValueToValueMapTy &VMap) {
+static void PropagateCallSiteMetadata(CallBase &CB, Function::iterator FStart,
+                                      Function::iterator FEnd) {
   MDNode *MemParallelLoopAccess =
       CB.getMetadata(LLVMContext::MD_mem_parallel_loop_access);
   MDNode *AccessGroup = CB.getMetadata(LLVMContext::MD_access_group);
   MDNode *AliasScope = CB.getMetadata(LLVMContext::MD_alias_scope);
   MDNode *NoAlias = CB.getMetadata(LLVMContext::MD_noalias);
   if (!MemParallelLoopAccess && !AccessGroup && !AliasScope && !NoAlias)
     return;
 
-  for (ValueToValueMapTy::iterator VMI = VMap.begin(), VMIE = VMap.end();
-       VMI != VMIE; ++VMI) {
-    // Check that key is an instruction, to skip the Argument mapping, which
-    // points to an instruction in the original function, not the inlined one.
-    if (!VMI->second || !isa<Instruction>(VMI->first))
-      continue;
-
-    Instruction *NI = dyn_cast<Instruction>(VMI->second);
-    if (!NI)
-      continue;
-
+  for (BasicBlock &BB : make_range(FStart, FEnd)) {
+    for (Instruction &I : BB) {
       // This metadata is only relevant for instructions that access memory.
-    if (!NI->mayReadOrWriteMemory())
+      if (!I.mayReadOrWriteMemory())
         continue;
 
       if (MemParallelLoopAccess) {
         // TODO: This probably should not overwrite MemParalleLoopAccess.
         MemParallelLoopAccess = MDNode::concatenate(
-          NI->getMetadata(LLVMContext::MD_mem_parallel_loop_access),
+            I.getMetadata(LLVMContext::MD_mem_parallel_loop_access),
             MemParallelLoopAccess);
-      NI->setMetadata(LLVMContext::MD_mem_parallel_loop_access,
+        I.setMetadata(LLVMContext::MD_mem_parallel_loop_access,
                       MemParallelLoopAccess);
       }
 
       if (AccessGroup)
-      NI->setMetadata(LLVMContext::MD_access_group, uniteAccessGroups(
-          NI->getMetadata(LLVMContext::MD_access_group), AccessGroup));
+        I.setMetadata(LLVMContext::MD_access_group, uniteAccessGroups(
+            I.getMetadata(LLVMContext::MD_access_group), AccessGroup));
 
       if (AliasScope)
-      NI->setMetadata(LLVMContext::MD_alias_scope, MDNode::concatenate(
-          NI->getMetadata(LLVMContext::MD_alias_scope), AliasScope));
+        I.setMetadata(LLVMContext::MD_alias_scope, MDNode::concatenate(
+            I.getMetadata(LLVMContext::MD_alias_scope), AliasScope));
 
       if (NoAlias)
-      NI->setMetadata(LLVMContext::MD_noalias, MDNode::concatenate(
-          NI->getMetadata(LLVMContext::MD_noalias), NoAlias));
+        I.setMetadata(LLVMContext::MD_noalias, MDNode::concatenate(
+            I.getMetadata(LLVMContext::MD_noalias), NoAlias));
+    }
   }
 }
 
 /// Utility for cloning !noalias and !alias.scope metadata. When a code region
 /// using scoped alias metadata is inlined, the aliasing relationships may not
 /// hold between the two version. It is necessary to create a deep clone of the
 /// metadata, putting the two versions in separate scope domains.
 class ScopedAliasMetadataDeepCloner {
   using MetadataMap = DenseMap<const MDNode *, TrackingMDNodeRef>;
   SetVector<const MDNode *> MD;
   MetadataMap MDMap;
   void addRecursiveMetadataUses();
 
 public:
   ScopedAliasMetadataDeepCloner(const Function *F);
 
   /// Create a new clone of the scoped alias metadata, which will be used by
   /// subsequent remap() calls.
   void clone();
 
-  /// Remap instructions in the given VMap from the original to the cloned
+  /// Remap instructions in the given range from the original to the cloned
   /// metadata.
-  void remap(ValueToValueMapTy &VMap);
+  void remap(Function::iterator FStart, Function::iterator FEnd);
 };
 
 ScopedAliasMetadataDeepCloner::ScopedAliasMetadataDeepCloner(
     const Function *F) {
   for (const BasicBlock &BB : *F) {
     for (const Instruction &I : BB) {
       if (const MDNode *M = I.getMetadata(LLVMContext::MD_alias_scope))
         MD.insert(M);
for (const MDNode *I : MD) {
     MDTuple *TempM = cast<MDTuple>(MDMap[I]);
     assert(TempM->isTemporary() && "Expected temporary node");
 
     TempM->replaceAllUsesWith(NewM);
     NewOps.clear();
   }
 }
 
-void ScopedAliasMetadataDeepCloner::remap(ValueToValueMapTy &VMap) {
+void ScopedAliasMetadataDeepCloner::remap(Function::iterator FStart,
+                                          Function::iterator FEnd) {
   if (MDMap.empty())
     return; // Nothing to do.
 
-  for (auto Entry : VMap) {
-    // Check that key is an instruction, to skip the Argument mapping, which
-    // points to an instruction in the original function, not the inlined one.
-    if (!Entry->second || !isa<Instruction>(Entry->first))
-      continue;
-
-    Instruction *I = dyn_cast<Instruction>(Entry->second);
-    if (!I)
-      continue;
-
-    // Only update scopes when we find them in the map. If they are not, it is
-    // because we already handled that instruction before. This is faster than
-    // tracking which instructions we already updated.
-    if (MDNode *M = I->getMetadata(LLVMContext::MD_alias_scope))
+  for (BasicBlock &BB : make_range(FStart, FEnd)) {
+    for (Instruction &I : BB) {
+      // TODO: The null checks for the MDMap.lookup() results should no longer
+      // be necessary.
+      if (MDNode *M = I.getMetadata(LLVMContext::MD_alias_scope))
         if (MDNode *MNew = MDMap.lookup(M))
-        I->setMetadata(LLVMContext::MD_alias_scope, MNew);
+          I.setMetadata(LLVMContext::MD_alias_scope, MNew);
 
-    if (MDNode *M = I->getMetadata(LLVMContext::MD_noalias))
+      if (MDNode *M = I.getMetadata(LLVMContext::MD_noalias))
         if (MDNode *MNew = MDMap.lookup(M))
-        I->setMetadata(LLVMContext::MD_noalias, MNew);
+          I.setMetadata(LLVMContext::MD_noalias, MNew);
 
-    if (auto *Decl = dyn_cast<NoAliasScopeDeclInst>(I))
+      if (auto *Decl = dyn_cast<NoAliasScopeDeclInst>(&I))
         if (MDNode *MNew = MDMap.lookup(Decl->getScopeList()))
           Decl->setScopeList(MNew);
     }
   }
+}
 
 /// If the inlined function has noalias arguments,
 /// then add new alias scopes for each noalias argument, tag the mapped noalias
 /// parameters with noalias metadata specifying the new scope, and tag all
 /// non-derived loads, stores and memory intrinsics with the new alias scopes.
 static void AddAliasScopeMetadata(CallBase &CB, ValueToValueMapTy &VMap,
                                   const DataLayout &DL, AAResults *CalleeAAR) {
   if (!EnableNoAliasConversion)
if (UseNoAliasIntrinsic) {
       // Ignore the result for now. The result will be used when the
       // llvm.noalias intrinsic is introduced.
       (void)NoAliasDecl;
     }
   }
 
   // Iterate over all new instructions in the map; for all memory-access
   // instructions, add the alias scope metadata.
   for (ValueToValueMapTy::iterator VMI = VMap.begin(), VMIE = VMap.end();
        VMI != VMIE; ++VMI) {
     if (const Instruction *I = dyn_cast<Instruction>(VMI->first)) {
       if (!VMI->second)
         continue;
 
       Instruction *NI = dyn_cast<Instruction>(VMI->second);
       if (!NI)
         continue;
Function::iterator FirstNewBlock;
     // For 'nodebug' functions, the associated DISubprogram is always null.
     // Conservatively avoid propagating the callsite debug location to
     // instructions inlined from a function whose DISubprogram is not null.
     fixupLineNumbers(Caller, FirstNewBlock, &CB,
                      CalledFunc->getSubprogram() != nullptr);
 
     // Now clone the inlined noalias scope metadata.
     SAMetadataCloner.clone();
-    SAMetadataCloner.remap(VMap);
+    SAMetadataCloner.remap(FirstNewBlock, Caller->end());
 
     // Add noalias metadata if necessary.
     AddAliasScopeMetadata(CB, VMap, DL, CalleeAAR);
 
     // Clone return attributes on the callsite into the calls within the inlined
     // function which feed into its return value.
     AddReturnAttributes(CB, VMap);
 
     // Propagate metadata on the callsite if necessary.
-    PropagateCallSiteMetadata(CB, VMap);
+    PropagateCallSiteMetadata(CB, FirstNewBlock, Caller->end());
 
     // Register any cloned assumptions.
     if (IFI.GetAssumptionCache)
       for (BasicBlock &NewBlock :
            make_range(FirstNewBlock->getIterator(), Caller->end()))
         for (Instruction &I : NewBlock)
           if (auto *II = dyn_cast<AssumeInst>(&I))
             IFI.GetAssumptionCache(*Caller).registerAssumption(II);

llvm/test/Transforms/Inline/pr50270.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -inline < %s | FileCheck %s
+
+; This tests cases where instructions in the callee are simplified to
+; instructions in the caller, thus making VMap contain instructions from
+; the caller. We should not be assigning incorrect noalias metadata in
+; that case.
+
+declare { i64* } @opaque_callee()
+
+define { i64* } @callee(i64* %x) {
+; CHECK-LABEL: @callee(
+; CHECK-NEXT:    [[RES:%.*]] = insertvalue { i64* } undef, i64* [[X:%.*]], 0
+; CHECK-NEXT:    ret { i64* } [[RES]]
+;
+  %res = insertvalue { i64* } undef, i64* %x, 0
+  ret { i64* } %res
+}
+
+; @opaque_callee() should not receive noalias metadata here.
+define void @caller() {
+; CHECK-LABEL: @caller(
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !0)
+; CHECK-NEXT:    [[S:%.*]] = call { i64* } @opaque_callee()
+; CHECK-NEXT:    [[X:%.*]] = extractvalue { i64* } [[S]], 0
+; CHECK-NEXT:    ret void
+;
+  call void @llvm.experimental.noalias.scope.decl(metadata !0)
+  %s = call { i64* } @opaque_callee()
+  %x = extractvalue { i64* } %s, 0
+  call { i64* } @callee(i64* %x), !noalias !0
+  ret void
+}
+
+; @opaque_callee() should no the same noalias metadata as the load from the
+; else branch, not as the load in the if branch.
+define { i64* } @self_caller(i1 %c, i64* %a) {
+; CHECK-LABEL: @self_caller(
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !0)
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    [[S:%.*]] = call { i64* } @opaque_callee(), !noalias !0
+; CHECK-NEXT:    [[X:%.*]] = extractvalue { i64* } [[S]], 0
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !3)
+; CHECK-NEXT:    [[TMP1:%.*]] = load volatile i64, i64* [[X]], align 4, !alias.scope !3
+; CHECK-NEXT:    ret { i64* } [[S]]
+; CHECK:       else:
+; CHECK-NEXT:    [[R2:%.*]] = insertvalue { i64* } undef, i64* [[A:%.*]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = load volatile i64, i64* [[A]], align 4, !alias.scope !0
+; CHECK-NEXT:    ret { i64* } [[R2]]
+;
+  call void @llvm.experimental.noalias.scope.decl(metadata !0)
+  br i1 %c, label %if, label %else
+
+if:
+  %s = call { i64* } @opaque_callee(), !noalias !0
+  %x = extractvalue { i64* } %s, 0
+  %r = call { i64* } @self_caller(i1 false, i64* %x)
+  ret { i64* } %r
+
+else:
+  %r2 = insertvalue { i64* } undef, i64* %a, 0
+  load volatile i64, i64* %a, !alias.scope !0
+  ret { i64* } %r2
+}
+
+declare void @llvm.experimental.noalias.scope.decl(metadata)
+
+!0 = !{!1}
+!1 = !{!1, !2, !"scope"}
+!2 = !{!2, !"domain"}