[IRMover] Remove UB implying parameter attributes when necessary

[TimNN]

Summary

When importing functions from some module X into some module Y, they may reference other functions already present in Y. The signature (especially attributes) of those other functions may have diverged between X and Y (e.g. due to the Dead Argument Elimination optimization). If necessary, modify the attributes to avoid UB.

See the added test and implementation comments for more details.

This was exposed by https://reviews.llvm.org/D133036 before it was reverted. Fixes https://github.com/llvm/llvm-project/issues/58976.

[tejohnson]

llvm/lib/Linker/IRMover.cpp

656

It seems this can/should only be an issue when the dest module only contains a declaration, and we are only importing a declaration. Should we be checking that F->isDeclaration() (or doing it after that check falls through)? And should we further check that we wouldn't invoke linkGlobalValueBody() further below? Wondering if the change to ensure function attr compatibility is either unnecessary or would be an error in either of those cases.

llvm/test/Transforms/FunctionImport/attr_fixup.ll

12

Suggest making Inputs/attr_fixup.ll contain the post DAE code, rather than running that and expecting it will make the necessary change.

[TimNN]

llvm/lib/Linker/IRMover.cpp

656

Doing this only when both are declarations makes sense, I think. (Either way, I'll move this to a separate if statement because all the cases here handle the "If we already created the body, just return." case).

While thinking about this I found an (IMO) more severe issue. I was initially thinking that definitions with internal linkage may need to get their attributes updated. Then I realized that reusing an existing symbol with internal linkage when importing a function was probably always wrong (because they could have completely different definitions).

When I experimented with this, I found that the current code actually _does_ (incorrectly, IIUC) reuse symbols with internal linkage, see https://github.com/llvm/llvm-project/issues/59347 for a repro. I feel like that should be fixed first before continuing this patch.

[tejohnson]

llvm/lib/Linker/IRMover.cpp

656

We promote to global scope and rename any local functions that are either exported by reference or definition to another module. I can see why that didn't happen with the test case you created, as you didn't do any post thin-link updates on the other module. Will make a comment on the github issue.

[TimNN]

Address review comments

• Move the fixup to after the existing if.
• Only perform the fixup when the source is a definition and the destination a declaration.
• Change the test to contain the post-deadargelim IR.

[TimNN]

Thanks for the review and also your input on https://github.com/llvm/llvm-project/issues/59347!

llvm/lib/Linker/IRMover.cpp

656

I'm not sure about linkGlobalValueBody(), I left things as-is with regard to that. Let me know if you want me to take this into account.

[TimNN]

Sync to HEAD to check if that resolves the (IMO likely unrelated) test timeout.

[tejohnson]

llvm/lib/Linker/IRMover.cpp

608

I thought this could only happen if Src's definition is not linked in.

Ah, Src may contain a def but we may not link in that def (depending on whether we decide to invoke linkGlobalValueBody). I believe if the def of Src was linked in, we wouldn't have an issue, since presumably that def would contain the correct attributes? I assume if you remove the "noinline" from inner in the Input file for the test it's def gets linked in - I forget, does that work correctly?

Assuming the answer to the last question above is yes, then I think you would only want to call this when we don't invoke linkGlobalValueBody.

Then I don't think you need or want this if statement at all (we've already checked in the caller that dest is a declaration).

[TimNN]

Only call ensureFunctionCompatibility if linkGlobalValueBody isn't called.

[TimNN]

llvm/lib/Linker/IRMover.cpp

608

Ack, makes sense, I think I understand this code a bit better now. Removed this if statement and change the code to only call ensureFunctionCompatibility if the body did not get linked.

And yes, removing the noinline fixes the problem (though the original issues happened without noinline, i.e. the optimizier decided to not inline @inner for some other reason).

[tejohnson]

Sorry for the delay, I had another review of attribute updates come in (D137360), can you work with @gulfem to have a consistent place to do these updates? I'm thinking that there could be a single "updateAttributes" function, and in there you could invoke your ensureFunctionCompatibility if New is still a declaration (see D137360, I've suggested removeAttributes be renamed to updateAttributes).

[gulfem]

In D139209#4008098, @tejohnson wrote:

Sorry for the delay, I had another review of attribute updates come in (D137360), can you work with @gulfem to have a consistent place to do these updates? I'm thinking that there could be a single "updateAttributes" function, and in there you could invoke your ensureFunctionCompatibility if New is still a declaration (see D137360, I've suggested removeAttributes be renamed to updateAttributes).

@TimNN,
Based on @tejohnson 's suggestion, I renamed the function to updateAttributes in D137360. How would you like to coordinate? I'm thinking that I can land D137360, and you can extend updateAttributes() function to include a src and dest arguments, and call your ensureFunctionCompatibility() function from there. Does that sound good or do you have another suggestion?

[TimNN]

In D139209#4013911, @gulfem wrote:

I'm thinking that I can land D137360, and you can extend updateAttributes() function to include a src and dest arguments, and call your ensureFunctionCompatibility() function from there.

Yeah, sounds good! I don't know when I'll have time to come back to this during the holidays, so go ahead and submit yours if is ready and I'll adapt.

[gulfem]

In D139209#4013982, @TimNN wrote:

In D139209#4013911, @gulfem wrote:

I'm thinking that I can land D137360, and you can extend updateAttributes() function to include a src and dest arguments, and call your ensureFunctionCompatibility() function from there.

Yeah, sounds good! I don't know when I'll have time to come back to this during the holidays, so go ahead and submit yours if is ready and I'll adapt.

Great, I landed D137360! Feel free to change updateAttributes() function when you revisit this. Happy holidays!

[TimNN]

Adapt to changes from D137360.

This reverts to not taking into account whether linkGlobalValueBody was called because that information is not readily available in the helper added by D137360 and I did not want to proxy that information via another parameter.

[arsenm]

llvm/lib/Linker/IRMover.cpp

1580–1581

Drive by comment, but is this brokenly not checking for CallBase?

llvm/test/Transforms/FunctionImport/Inputs/attr_fixup_dae_noundef.ll

6 ↗

(On Diff #488941)

Shouldn't use anonymous values in tests

[arsenm]

llvm/lib/Linker/IRMover.cpp

1594–1596

I thought we casted the function type at callsites for different signatures in different modules, so this wouldn't hold?

[TimNN]

Don't use anonymous values in tests

[TimNN]

llvm/lib/Linker/IRMover.cpp

1580–1581

I'm not familiary enough with this code / LLVM to answer this. cc @gulfem who added this in D137360.

Anyway, marking this comment as "done" because if this needs to be fixed, it shouldn't be part of this change.

1594–1596

@tejohnson do you have any input on this?

(I last looked at this code in-depth in early December, but don't remember seeing anything relevant. ninja check-all is also clean and in a brief test IRMover would happily import a function with a different signature without creating any bitcasts).

[arsenm]

llvm/lib/Linker/IRMover.cpp

1594–1596

These days it would appear as a callsite with a different type from the callee, there's no explicit bitcast with opaque pointers

[TimNN]

llvm/lib/Linker/IRMover.cpp

1594–1596

Ack, then I guess the "brief test" wasn't really useful.

Anyway, if it is indeed expected that SrcF and DstF have different signatures here, I'd appreciate some guidance on how to best detect and handle that.

[tejohnson]

llvm/lib/Linker/IRMover.cpp

1580–1581

Thanks for noting that, it is a limitation with the earlier patch that I missed. I will copy your comment over there.

1580–1581

How about just do an early return if !DstF, since your changes below also rely on DstF being non-null.

1585

I think this is backwards? My understanding (also from the code and the test case) is that the issue occurs if F is a declaration in the dest module but was a definition in the source module.

1594–1596

Sorry I'm not really sure. @arsenm do you mean that the arg_size could be different? Do you have an example of how/when that would happen?

[TimNN]

Fix definition/declaration wording in comment.

llvm/lib/Linker/IRMover.cpp

1580–1581

Given the generic name ("updateAttributes") and purpose of this function, I decided to avoid any early returns.

(Put differently: This function currently has two mostly independent (and long) sections, each with their own prerequisites. I wanted to keep the control-flow contained to the individual sections, especially in light of the possibility that this function may handle non-function globals in the future).

I'd be happy to add the early return, but in that case I'd advocate for renaming the function "updateFunctionAttributes".

1585

Yes, of course. If I don't pay very close attention it's too easy to mix up the two. I've updated the patch.

[tejohnson]

lgtm, but there is still an outstanding question to @arsenm. I'm not exactly sure of the scenario they are concerned about.

[TimNN]

llvm/lib/Linker/IRMover.cpp

1594–1596

Friendly ping, @arsenm! Could you provide more details about your concern, or even better, a concrete example?

[aeubanks]

llvm/lib/Linker/IRMover.cpp

1594–1596

I think what @arsenm is talking about is a separate thing that isn't relevant here. The callsite thing is on function calls (CallBase::getCalledFunction() returns nullptr if the call args and callee are mismatched), but this is about handling function declarations/definitions of the same function.

If there's any updating of callers of a function, that would be happening somewhere else, not in the code touched by this patch.

[TimNN]

llvm/lib/Linker/IRMover.cpp

1594–1596

Than you for your input, @aeubanks!

@tejohnson: Would you be fine with submitting this patch as-is?

[tejohnson]

llvm/lib/Linker/IRMover.cpp

1594–1596

Yes, lgtm

[TimNN]

Rebase onto HEAD.

[aeubanks]

I believe we're seeing a stage 2 ThinLTO build of clang crash due to this patch, will reduce test case

[aeubanks]

In D139209#4144952, @aeubanks wrote:

I believe we're seeing a stage 2 ThinLTO build of clang crash due to this patch, will reduce test case

confirmed locally, should be reproducible everywhere (saw the same crash on a windows bot)

reverting

[TimNN]

The current code is failing for a situation like this:

--- SRC ---
; Module Name: lib/libLLVMCodeGen.a(TargetLoweringBase.cpp.o at 275472348)
; Module Source: /home/logic/Projects/llvm/llvm/lib/CodeGen/TargetLoweringBase.cpp
; Materializable
; Function Attrs: mustprogress noinline nounwind optnone uwtable
define { ptr, i8 } @_ZNK4llvm18TargetLoweringBase23findRepresentativeClassEPKNS_18TargetRegisterInfoENS_3MVTE(ptr noundef nonnull align 8 dereferenceable(218195) %0, ptr noundef %1, i8 %2) unnamed_addr #1 align 2 {}

--- DST ---
; Module Name: lib/libLLVMAArch64CodeGen.a(AArch64ISelLowering.cpp.o at 48495784)
; Module Source: /home/logic/Projects/llvm/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
declare void @_ZNK4llvm18TargetLoweringBase23findRepresentativeClassEPKNS_18TargetRegisterInfoENS_3MVTE() unnamed_addr

It seems like when generating vtables, the code at https://github.com/llvm/llvm-project/blob/main/clang/lib/CodeGen/CGVTables.cpp may generate function declarations with an "empty/void" signature for the IsUnprototyped case.

Since clang is deliberately generating this code, I'm assuming this is considered valid according to LLVM IR semantics. I'm not sure what the most appropriate fix here would be. I primarily see three open questions:

Question 1: Is my assumption correct and this is valid LLVM IR?

Question 2: Should we attempt to detect this case specifically (DST has zero args and returns void) but otherwise still assert that the argument lists have the same length? Or should we handle all cases of argument lists with different lengths?

Question 3: For the cases detected by (1), should we conservatively remove UB-implying attributes from all arguments (erring on the side of correctness)? Or should we just bail out and not do anything (erring on the side of performance)?

[tejohnson]

In D139209#4146484, @TimNN wrote:

The current code is failing for a situation like this:

--- SRC ---
; Module Name: lib/libLLVMCodeGen.a(TargetLoweringBase.cpp.o at 275472348)
; Module Source: /home/logic/Projects/llvm/llvm/lib/CodeGen/TargetLoweringBase.cpp
; Materializable
; Function Attrs: mustprogress noinline nounwind optnone uwtable
define { ptr, i8 } @_ZNK4llvm18TargetLoweringBase23findRepresentativeClassEPKNS_18TargetRegisterInfoENS_3MVTE(ptr noundef nonnull align 8 dereferenceable(218195) %0, ptr noundef %1, i8 %2) unnamed_addr #1 align 2 {}

--- DST ---
; Module Name: lib/libLLVMAArch64CodeGen.a(AArch64ISelLowering.cpp.o at 48495784)
; Module Source: /home/logic/Projects/llvm/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
declare void @_ZNK4llvm18TargetLoweringBase23findRepresentativeClassEPKNS_18TargetRegisterInfoENS_3MVTE() unnamed_addr

It seems like when generating vtables, the code at https://github.com/llvm/llvm-project/blob/main/clang/lib/CodeGen/CGVTables.cpp may generate function declarations with an "empty/void" signature for the IsUnprototyped case.

Since clang is deliberately generating this code, I'm assuming this is considered valid according to LLVM IR semantics. I'm not sure what the most appropriate fix here would be. I primarily see three open questions:

Question 1: Is my assumption correct and this is valid LLVM IR?

Yes I think we should assume so. I believe there may be other cases where the compiler may generate a function declaration without knowing the exact definition's arguments...I know WPD does this sometimes although it would be after importing so wouldn't trigger this particular code.

Question 2: Should we attempt to detect this case specifically (DST has zero args and returns void) but otherwise still assert that the argument lists have the same length? Or should we handle all cases of argument lists with different lengths?

I think it might be best to handle these cases conservatively, eg:

Question 3: For the cases detected by (1), should we conservatively remove UB-implying attributes from all arguments (erring on the side of correctness)? Or should we just bail out and not do anything (erring on the side of performance)?

I think the best option is likely to remove all UB implying attributes from the dst in that case. Note that in the example shown above, it should not need to do anything, since the dst is a declaration without any args, which I would guess would be the likely case when we are going to hit this situation.

llvm/include/llvm/IR/Attributes.h

public:
   bool hasAttributes() const { return SetNode != nullptr; }
 
   /// Return true if the attribute exists in this set.
   bool hasAttribute(Attribute::AttrKind Kind) const;
 
   /// Return true if the attribute exists in this set.
   bool hasAttribute(StringRef Kind) const;
 
+  /// Return true if the attribute exists in this set.
+  bool hasAttribute(Attribute A) const {
+    if (A.isStringAttribute())
+      return hasAttribute(A.getKindAsString());
+    return hasAttribute(A.getKindAsEnum());
+  }
+
   /// Return the attribute object.
   Attribute getAttribute(Attribute::AttrKind Kind) const;
 
   /// Return the target-dependent attribute object.
   Attribute getAttribute(StringRef Kind) const;
 
   MaybeAlign getAlignment() const;
   MaybeAlign getStackAlignment() const;

llvm/lib/Linker/IRMover.cpp

class IRLinker {
   void flushRAUWWorklist();
 
   /// When importing for ThinLTO, prevent importing of types listed on
   /// the DICompileUnit that we don't need a copy of in the importing
   /// module.
   void prepareCompileUnitsForImport();
   void linkNamedMDNodes();
 
+  /// Ensures that the signature of the function in the destination module is
+  /// compatible with the signature of the function in the source module.
+  ///
+  /// `Dst` and `Src` should refer to the same function, possibly
+  /// declared/defined in different modules.
+  void ensureFunctionCompatibility(Function &Dst, Function &Src);
+
 public:
   IRLinker(Module &DstM, MDMapT &SharedMDs,
            IRMover::IdentifiedStructTypeSet &Set, std::unique_ptr<Module> SrcM,
            ArrayRef<GlobalValue *> ValuesToLink,
            IRMover::LazyCallback AddLazyFor, bool IsPerformingImport)
       : DstM(DstM), SrcM(std::move(SrcM)), AddLazyFor(std::move(AddLazyFor)),
         TypeMap(Set), GValMaterializer(*this), LValMaterializer(*this),
         SharedMDs(SharedMDs), IsPerformingImport(IsPerformingImport),

 Value *GlobalValueMaterializer::materialize(Value *SGV) {
   return TheIRLinker.materialize(SGV, false);
 }
 
 Value *LocalValueMaterializer::materialize(Value *SGV) {
   return TheIRLinker.materialize(SGV, true);
 }
 
+void IRLinker::ensureFunctionCompatibility(Function &Dst, Function &Src) {
+  // At this point, `Dst` and `Src` should refer to the same function, but
+  // declared/defined in different modules. All references to `Src` should get
+  // replaced by references to `Dst`.
+  //
+  // `Dst` and `Src` should generally have the same signature, but in some cases
+  // it is possible for the signatures (especially attributes) to diverge.
+  //
+  // For example, "dead argument elimination" may have removed a `noundef` from
+  // one of `Src`'s arguments. When replacing `Src` with `Dst`, we need to also
+  // remove the `noundef` from `Dst` because some other function imported from
+  // `Src`'s module may be relying on `noundef` _not_ being present.
+  //
+  // See the test in `FunctionImport/attr_fixup.ll` for a concrete example.
+  //
+  // The code below handles _known_ incompatibilities. Others may exist.
+
+  assert(Dst.arg_size() == Src.arg_size() &&
+         "Dst and Src should have the same signature.");
+
+  // Remove UB implying attributes present on `Dst`'s arguments but not `Src`'s.
+  AttributeMask UBImplyingAttributes =
+      AttributeFuncs::getUBImplyingAttributes();
+  for (size_t ArgI = 0; ArgI < Dst.arg_size(); ArgI++) {
+    AttributeSet DstAttrs = Dst.getAttributes().getParamAttrs(ArgI);
+    AttributeSet SrcAttrs = Src.getAttributes().getParamAttrs(ArgI);
+    AttributeMask ToRemove;
+
+    for (auto &Attr : DstAttrs) {
+      if (SrcAttrs.hasAttribute(Attr))
+        continue;
+
+      if (UBImplyingAttributes.contains(Attr))
+        ToRemove.addAttribute(Attr);
+    }
+
+    Dst.removeParamAttrs(ArgI, ToRemove);
+  }
+}
+
 Value *IRLinker::materialize(Value *V, bool ForIndirectSymbol) {
   auto *SGV = dyn_cast<GlobalValue>(V);
   if (!SGV)
     return nullptr;
 
   // When linking a global from other modules than source & dest, skip
   // materializing it because it would be mapped later when its containing
   // module is linked. Linking it now would potentially pull in many types that
if (!*NewProto)
     return nullptr;
 
   GlobalValue *New = dyn_cast<GlobalValue>(*NewProto);
   if (!New)
     return *NewProto;
 
   // If we already created the body, just return.
   if (auto *F = dyn_cast<Function>(New)) {
     if (!F->isDeclaration())
       return New;
   } else if (auto *V = dyn_cast<GlobalVariable>(New)) {
     if (V->hasInitializer() || V->hasAppendingLinkage())
       return New;
   } else if (auto *GA = dyn_cast<GlobalAlias>(New)) {
     if (GA->getAliasee())
       return New;
   } else if (auto *GI = dyn_cast<GlobalIFunc>(New)) {
Value *IRLinker::materialize(Value *V, bool ForIndirectSymbol) {
   // confirming the same value has been scheduled.  If there is an entry in the
   // ValueMap but the value is different, it means that the value already had a
   // definition in the destination module (linkonce for instance), but we need a
   // new definition for the indirect symbol ("New" will be different).
   if ((ForIndirectSymbol && ValueMap.lookup(SGV) == New) ||
       (!ForIndirectSymbol && IndirectSymbolValueMap.lookup(SGV) == New))
     return New;
 
-  if (ForIndirectSymbol || shouldLink(New, *SGV))
+  if (ForIndirectSymbol || shouldLink(New, *SGV)) {
     setError(linkGlobalValueBody(*New, *SGV));
+  } else {
+    // If the body did not get linked, some attributes may need to be updated.
+    if (auto *DF = dyn_cast<Function>(New)) {
+      if (auto *SF = dyn_cast<Function>(SGV)) {
+        ensureFunctionCompatibility(*DF, *SF);
+      }
+    }
+  }
 
   return New;
 }
 
 /// Loop through the global variables in the src module and merge them into the
 /// dest module.
 GlobalVariable *IRLinker::copyGlobalVariableProto(const GlobalVariable *SGVar) {
   // No linking to be performed or linking from the source: simply create an

 static std::string adjustInlineAsm(const std::string &InlineAsm,
                                    const Triple &Triple) {
   if (Triple.getArch() == Triple::thumb || Triple.getArch() == Triple::thumbeb)
     return ".text\n.balign 2\n.thumb\n" + InlineAsm;
   if (Triple.getArch() == Triple::arm || Triple.getArch() == Triple::armeb)
     return ".text\n.balign 4\n.arm\n" + InlineAsm;
   return InlineAsm;
 }
 
 Error IRLinker::run() {
   // Ensure metadata materialized before value mapping.
   if (SrcM->getMaterializer())
     if (Error Err = SrcM->getMaterializer()->materializeMetadata())
       return Err;
 
   // Inherit the target data from the source module if the destination module
   // doesn't have one already.
   if (DstM.getDataLayout().isDefault())
     DstM.setDataLayout(SrcM->getDataLayout());
 
   // Copy the target triple from the source to dest if the dest's is empty.
   if (DstM.getTargetTriple().empty() && !SrcM->getTargetTriple().empty())
     DstM.setTargetTriple(SrcM->getTargetTriple());
 
   Triple SrcTriple(SrcM->getTargetTriple()), DstTriple(DstM.getTargetTriple());
 
   // During CUDA compilation we have to link with the bitcode supplied with
   // CUDA. libdevice bitcode either has no data layout set (pre-CUDA-11), or has
   // the layout that is different from the one used by LLVM/clang (it does not
   // include i128). Issuing a warning is not very helpful as there's not much
   // the user can do about it.
   bool EnableDLWarning = true;
   bool EnableTripleWarning = true;

llvm/test/Transforms/FunctionImport/Inputs/attr_fixup.ll

+; This file contains the post-"deadargelim" IR.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @outer(i32 noundef %0) {
+  ; The parameter was originally `noundef %0`, changed to `poison` by "deadargelim".
+  call void @inner(i32 poison)
+  ret void
+}
+
+; %0 was originally `noundef`, removed by "deadargelim".
+define void @inner(i32 %0) #0 {
+  ret void
+}
+
+attributes #0 = { noinline }

llvm/test/Transforms/FunctionImport/attr_fixup.ll

+; Test to ensure that if a definition is imported, already-present declarations
+; are updated as necessary: Definitions from the same module may be optimized
+; together. Thus care must be taken when importing only a subset of the
+; definitions from a module (because other referenced definitions from that
+; module may have been changed by the optimizer and may no longer match
+; declarations already present in the module being imported into).
+
+; Generate bitcode and index, and run the function import.
+; `Inputs/attr_fixup.ll` contains the post-"Dead Argument Elimination" IR, which
+; removed the `noundef` from `@inner`.
+; RUN: opt -module-summary %p/Inputs/attr_fixup.ll -o %t.inputs.attr_fixup.bc
+; RUN: opt -module-summary %s -o %t.main.bc
+; RUN: llvm-lto -thinlto -o %t.summary %t.main.bc %t.inputs.attr_fixup.bc
+; RUN: opt -passes=function-import -summary-file %t.summary.thinlto.bc %t.main.bc -S 2>&1 \
+; RUN:   | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @main()  {
+  call void @outer(i32 noundef 1)
+  call void @inner(i32 noundef 1)
+  ret void
+}
+
+; Because `@inner` is `noinline`, it should not get imported. However, the
+; `noundef` should be removed.
+; CHECK: declare void @inner(i32)
+declare void @inner(i32 noundef)
+
+; `@outer` should get imported.
+; CHECK: define available_externally void @outer(i32 noundef %0)
+; CHECK-NEXT: call void @inner(i32 poison)
+declare void @outer(i32 noundef)