diff --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h --- a/llvm/include/llvm/IR/Attributes.h +++ b/llvm/include/llvm/IR/Attributes.h @@ -359,6 +359,13 @@ /// 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; diff --git a/llvm/lib/Linker/IRMover.cpp b/llvm/lib/Linker/IRMover.cpp --- a/llvm/lib/Linker/IRMover.cpp +++ b/llvm/lib/Linker/IRMover.cpp @@ -526,6 +526,13 @@ 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 SrcM, @@ -580,6 +587,46 @@ 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(V); if (!SGV) @@ -606,6 +653,10 @@ // If we already created the body, just return. if (auto *F = dyn_cast(New)) { + if (auto *SF = dyn_cast(SGV)) { + ensureFunctionCompatibility(*F, *SF); + } + if (!F->isDeclaration()) return New; } else if (auto *V = dyn_cast(New)) { diff --git a/llvm/test/Transforms/FunctionImport/Inputs/attr_fixup.ll b/llvm/test/Transforms/FunctionImport/Inputs/attr_fixup.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/FunctionImport/Inputs/attr_fixup.ll @@ -0,0 +1,13 @@ +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) { + call void @inner(i32 noundef %0) + ret void +} + +define void @inner(i32 noundef %0) #0 { + ret void +} + +attributes #0 = { noinline } diff --git a/llvm/test/Transforms/FunctionImport/attr_fixup.ll b/llvm/test/Transforms/FunctionImport/attr_fixup.ll new file mode 100644 --- /dev/null +++ b/llvm/test/Transforms/FunctionImport/attr_fixup.ll @@ -0,0 +1,33 @@ +; 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 for the definitions, and run Dead Argument Elimination on +; them. This makes `@outer` call `@inner` with `poison` as the argument, while +; also removing `noundef` from `@inner`. +; RUN: opt -module-summary -passes=deadargelim %p/Inputs/attr_fixup.ll -o %t.inputs.attr_fixup.bc + +; Now generate the remaining bitcode and index, and run the function import. +; 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 + +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)