[IRMover] Remove UB implying parameter attributes when necessary

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 {
 
   /// 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();
 
   ///  Update attributes while linking.
-  void updateAttributes(GlobalValue &GV);
+  void updateAttributes(GlobalValue &DstGV, GlobalValue &SrcGV);
 
 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),
Value *IRLinker::materialize(Value *V, bool ForIndirectSymbol) {
   // 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))
     setError(linkGlobalValueBody(*New, *SGV));
 
-  updateAttributes(*New);
+  updateAttributes(*New, *SGV);
   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
   // identical version of the symbol over in the dest module... the
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;
 }
 
-void IRLinker::updateAttributes(GlobalValue &GV) {
+void IRLinker::updateAttributes(GlobalValue &DstGV, GlobalValue &SrcGV) {
+  auto *DstF = dyn_cast<Function>(&DstGV);
+  auto *SrcF = dyn_cast<Function>(&SrcGV);
+
   /// Remove nocallback attribute while linking, because nocallback attribute
   /// indicates that the function is only allowed to jump back into caller's
   /// module only by a return or an exception. When modules are linked, this
   /// property cannot be guaranteed anymore. For example, the nocallback
   /// function may contain a call to another module. But if we merge its caller
   /// and callee module here, and not the module containing the nocallback
   /// function definition itself, the nocallback property will be violated
   /// (since the nocallback function will call back into the newly merged module
   /// containing both its caller and callee). This could happen if the module
   /// containing the nocallback function definition is native code, so it does
   /// not participate in the LTO link. Note if the nocallback function does
   /// participate in the LTO link, and thus ends up in the merged module
   /// containing its caller and callee, removing the attribute doesn't hurt as
   /// it has no effect on definitions in the same module.
-  if (auto *F = dyn_cast<Function>(&GV)) {
-    if (!F->isIntrinsic())
-      F->removeFnAttr(llvm::Attribute::NoCallback);
+  if (DstF) {
+    if (!DstF->isIntrinsic())
+      DstF->removeFnAttr(llvm::Attribute::NoCallback);
 
     // Remove nocallback attribute when it is on a call-site.
-    for (BasicBlock &BB : *F)
+    for (BasicBlock &BB : *DstF)
       for (Instruction &I : BB)
         if (CallInst *CI = dyn_cast<CallInst>(&I))
           CI->removeFnAttr(Attribute::NoCallback);
   }
+
+  // Fix compatibility of diverged function signatures.
+  //
+  // If `F` was a declaration in its source module but is (already) a
+  // definition in the destination module, then the signatures of the two
+  // functions may have diverged (even if they were originally idential) due to
+  // optimizations performed on the source module.
+  //
+  // See the test in `FunctionImport/attr_fixup_dae_noundef.ll` for an example
+  // based on how Dead Argument Elimination can remove `noundef` from function
+  // arguments.
+  //
+  // The code below handles _known_ incompatibilities. Others may exist.
+  if (DstF && SrcF && DstF->isDeclaration() && !SrcF->isDeclaration()) {
+    assert(DstF->arg_size() == SrcF->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 < DstF->arg_size(); ArgI++) {
+      AttributeSet DstAttrs = DstF->getAttributes().getParamAttrs(ArgI);
+      AttributeSet SrcAttrs = SrcF->getAttributes().getParamAttrs(ArgI);
+      AttributeMask ToRemove;
+
+      for (auto &Attr : DstAttrs) {
+        if (SrcAttrs.hasAttribute(Attr))
+          continue;
+
+        if (UBImplyingAttributes.contains(Attr))
+          ToRemove.addAttribute(Attr);
+      }
+
+      DstF->removeParamAttrs(ArgI, ToRemove);
+    }
+  }
 }
 
 Error IRLinker::run() {
   // Ensure metadata materialized before value mapping.
   if (SrcM->getMaterializer())
     if (Error Err = SrcM->getMaterializer()->materializeMetadata())
       return Err;
 

llvm/test/Transforms/FunctionImport/Inputs/attr_fixup_dae_noundef.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_dae_noundef.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_dae_noundef.ll` contains the post-"Dead Argument Elimination" IR, which
+; removed the `noundef` from `@inner`.
+; RUN: opt -module-summary %p/Inputs/attr_fixup_dae_noundef.ll -o %t.inputs.attr_fixup_dae_noundef.bc
+; RUN: opt -module-summary %s -o %t.main.bc
+; RUN: llvm-lto -thinlto -o %t.summary %t.main.bc %t.inputs.attr_fixup_dae_noundef.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)