[Inliner] Update !associated metadata during inlining

[phosek]

Summary

When inlining a function that has !associated metadata, update the
metadata to point to caller instead of callee.

[phosek]

This is a prerequisite for D76802 because InstrProfiling pass runs before inliner which can invalidate the metadata.

[eugenis]

LGTM

llvm/lib/Transforms/Utils/InlineFunction.cpp

1148

This is technically N^2 if you do this for every function. I guess it's fine as long as you don't see a bump in the compilation time.

[phosek]

llvm/lib/Transforms/Utils/InlineFunction.cpp

1148

I don't know if the N^2 behavior is avoidable given the structure of the data where the !associated metadata is attached to GlobalValues and AFAIK there's not way to enumerate all values use a particular MDNode. I've at least added extra check to ensure that we only hit this if the value has metadata attached.

[eugenis]

N^2 could be avoided by building a function->associated global_value map to be used in all inlining in the current module, but that seems difficult to plumb through the layers, especially given that the inliner is an SCC pass, not a module pass.
Probably not worth it unless there is data that it slows down compilation time.

[phosek]

So after testing this change a bit more, I found a fundamental issue which is the handling of cases when a function is inlined into multiple functions.

Say we have the following IR:

@__inner_internal_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0

define internal void @inner_internal() {
  %1 = load i64, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 0, i64 1), align 8
  %2 = add i64 %1, 1
  store i64 %2, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 0, i64 1), align 8
  ret void
}

define void @outer1() {
  call void @inner_internal()
  ret void
}

define void @outer2() {
  call void @inner_internal()
  ret void
}

!0 = !{void ()* @inner_internal}

Without this patch, we end up with:

@__inner_internal_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0

define void @outer1() {
  %1 = load i64, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  %2 = add i64 %1, 1
  store i64 %2, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  ret void
}

define void @outer2() {
  %1 = load i64, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  %2 = add i64 %1, 1
  store i64 %2, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  ret void
}

!0 = distinct !{null}

With this patch, we end up with:

@__inner_internal_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0

define void @outer1() {
  %1 = load i64, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  %2 = add i64 %1, 1
  store i64 %2, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  ret void
}

define void @outer2() {
  %1 = load i64, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  %2 = add i64 %1, 1
  store i64 %2, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @__inner_internal_metadata, i64 1, i64 0), align 8
  ret void
}

!0 = !{void ()* @outer2}

Both of these are wrong. In the first case we end up with broken metadata, in the second case the metadata is pointing only to a single function where it should be pointing to both. I'm not sure how to solve this, since there isn't a way to express this with SHF_LINK_ORDER since a single section could only be linked to a single section, but in this case we need a single section (with metadata) to be linked to multiple sections. Any thoughts?

[phosek]

One option we've discussed internally within our team would be to use ELF section groups (not COMDAT which is a special case of section groups) instead of SHF_LINK_ORDER, but those aren't supported in LLVM IR at the moment, and it's also counter to what was discussed in https://groups.google.com/g/generic-abi/c/_CbBM6T6WeM/m/LZnqx1KZAQAJ.

[pcc]

In this case you might consider replacing the global reference with null and implementing my proposal in https://bugs.llvm.org/show_bug.cgi?id=41734#c2 to make the section SHF_LINK_ORDER with a zero-length .init_array section.

[phosek]

In D80186#2067107, @pcc wrote:

In this case you might consider replacing the global reference with null and implementing my proposal in https://bugs.llvm.org/show_bug.cgi?id=41734#c2 to make the section SHF_LINK_ORDER with a zero-length .init_array section.

Wouldn't that mean that symbols associated with null won't be GCed by the linker?

[pcc]

Yeah. Thinking about it more, you may consider a different !associated scheme for the __prof*_ globals.

• __profc_* and __profn_* are associated with a dummy global (to allow the sections to be GC'd).
• __profd_* is associated with the corresponding __profc_* section.

Now no special rules are required for inlining. __profc_* is kept alive via the references in the functions, __profd_ is kept alive via !associated and __profn_* is kept alive by direct reference.

[phosek]

In D80186#2067298, @pcc wrote:

Yeah. Thinking about it more, you may consider a different !associated scheme for the __prof*_ globals.

• __profc_* and __profn_* are associated with a dummy global (to allow the sections to be GC'd).
• __profd_* is associated with the corresponding __profc_* section.

Now no special rules are required for inlining. __profc_* is kept alive via the references in the functions, __profd_ is kept alive via !associated and __profn_* is kept alive by direct reference.

I tried that, but now we get an error ld.lld: error: /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_cnts): sh_link points to discarded section /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_dummy) because the dummy global was discarded which creates a broken sh_link.

[pcc]

In D80186#2068055, @phosek wrote:

In D80186#2067298, @pcc wrote:

Yeah. Thinking about it more, you may consider a different !associated scheme for the __prof*_ globals.

• __profc_* and __profn_* are associated with a dummy global (to allow the sections to be GC'd).
• __profd_* is associated with the corresponding __profc_* section.

Now no special rules are required for inlining. __profc_* is kept alive via the references in the functions, __profd_ is kept alive via !associated and __profn_* is kept alive by direct reference.

I tried that, but now we get an error ld.lld: error: /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_cnts): sh_link points to discarded section /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_dummy) because the dummy global was discarded which creates a broken sh_link.

Hmm. Maybe instead of making the sections sh_link to a dummy global, make them sh_link to themselves. There was a proposal for sh_link=0 to mean what we want here (D72904) but I believe that sh_link=self should have the same semantics.

[phosek]

In D80186#2069110, @pcc wrote:

In D80186#2068055, @phosek wrote:

In D80186#2067298, @pcc wrote:

Yeah. Thinking about it more, you may consider a different !associated scheme for the __prof*_ globals.

• __profc_* and __profn_* are associated with a dummy global (to allow the sections to be GC'd).
• __profd_* is associated with the corresponding __profc_* section.

Now no special rules are required for inlining. __profc_* is kept alive via the references in the functions, __profd_ is kept alive via !associated and __profn_* is kept alive by direct reference.

I tried that, but now we get an error ld.lld: error: /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_cnts): sh_link points to discarded section /tmp/instrprof-gc-sections-5f007f.o:(__llvm_prf_dummy) because the dummy global was discarded which creates a broken sh_link.

Hmm. Maybe instead of making the sections sh_link to a dummy global, make them sh_link to themselves. There was a proposal for sh_link=0 to mean what we want here (D72904) but I believe that sh_link=self should have the same semantics.

Thanks, that seems to work!

[phosek]

This is no longer needed for the profile instrumentation.

[davidxl]

I saw you removed the dependency on another patch which was abandoned. Why is it no longer needed?

Also can you collect data of instrumented Clang with IR PGO:

1. the instrumented binary size change with/without this patch
2. the raw profile data size change with/without the patch.

llvm/lib/Transforms/Utils/InlineFunction.cpp

if (const Instruction *I = dyn_cast<Instruction>(VMI->first)) {
         NI->setMetadata(
             LLVMContext::MD_alias_scope,
             MDNode::concatenate(NI->getMetadata(LLVMContext::MD_alias_scope),
                                 MDNode::get(CalledFunc->getContext(), Scopes)));
     }
   }
 }
 
+static void UpdateAssociatedMetadata(CallBase &CB) {
+  Module *M = CB.getModule();
+
+  Function *Callee = CB.getCalledFunction();
+  if (!Callee->isUsedByMetadata())
+    return;
+  ValueAsMetadata *CalleeMetadata = ValueAsMetadata::getIfExists(Callee);
+  assert(CalleeMetadata && "Expected existing metadata");
+  MDNode *CalleeMDNode = MDNode::get(CB.getContext(), CalleeMetadata);
+
+  Function *Caller = CB.getFunction();
+  ValueAsMetadata *CallerMetadata = ValueAsMetadata::get(Caller);
+  MDNode *CallerMDNode = MDNode::get(CB.getContext(), CallerMetadata);
+
+  for (GlobalVariable &GV : M->globals())
+    if (GV.getMetadata(LLVMContext::MD_associated) == CalleeMDNode)
+      GV.setMetadata(LLVMContext::MD_associated, CallerMDNode);
+}
+
 static bool MayContainThrowingOrExitingCall(Instruction *Begin,
                                             Instruction *End) {
 
   assert(Begin->getParent() == End->getParent() &&
          "Expected to be in same basic block!");
   unsigned NumInstChecked = 0;
   // Check that all instructions in the range [Begin, End) are guaranteed to
   // transfer execution to successor.
if (IFI.GetAssumptionCache)
            make_range(FirstNewBlock->getIterator(), Caller->end()))
         for (Instruction &I : NewBlock) {
           if (auto *II = dyn_cast<IntrinsicInst>(&I))
             if (II->getIntrinsicID() == Intrinsic::assume)
               (*IFI.GetAssumptionCache)(*Caller).registerAssumption(II);
         }
   }
 
+  // Adjust !associated metadata to point to caller instead of callee.
+  UpdateAssociatedMetadata(CB);
+
   // If there are any alloca instructions in the block that used to be the entry
   // block for the callee, move them to the entry block of the caller.  First
   // calculate which instruction they should be inserted before.  We insert the
   // instructions at the end of the current alloca list.
   {
     BasicBlock::iterator InsertPoint = Caller->begin()->begin();
     for (BasicBlock::iterator I = FirstNewBlock->begin(),
          E = FirstNewBlock->end(); I != E; ) {

llvm/test/Transforms/Inline/inline-md-associated.ll

+; RUN: opt < %s -inline --inline-threshold=0 -S | FileCheck %s
+
+; Test that the inliner updates the !associated metadata for inlined functions.
+
+; CHECK: @__foo_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0
+@__foo_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0
+; CHECK: @__bar_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !0
+@__bar_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !1
+; CHECK: @__baz_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !1
+@__baz_metadata = private global [1 x i64] zeroinitializer, align 8, !associated !2
+
+define void @foo() {
+  call void @bar(i1 true)
+  ret void
+}
+
+define internal void @bar(i1 %p) {
+  br i1 %p, label %bb1, label %bb2
+
+bb1:
+  call void @foo()
+  ret void
+
+bb2:
+  call void @bar(i1 true)
+  ret void
+}
+
+define void @baz() {
+; CHECK: call void @foo()
+  call void @bar(i1 true)
+; CHECK: call void @foo()
+  call void @bar(i1 false)
+  ret void
+}
+
+; CHECK: !0 = !{void ()* @foo}
+!0 = !{void ()* @foo}
+!1 = !{void (i1)* @bar}
+; CHECK: !1 = !{void ()* @baz}
+!2 = !{void ()* @baz}