IR: Fix issues with linking typed sret

[arsenm]

Summary

Fixes issues causing the revert in
eb9f7c28e5fe6d75fed3587023e17f2997c8024b.

[dblaikie]

The goal is to recommit the old patch, with these changes in addition in the same commit? (or can this be committed first, then the other patch recommitted on top of this safely?)

[arsenm]

In D88973#2317661, @dblaikie wrote:

The goal is to recommit the old patch, with these changes in addition in the same commit? (or can this be committed first, then the other patch recommitted on top of this safely?)

I was going to squash this into the un-revert, but I could also split enough of the Attribute utility functions to pre-commit this

[rnk]

Fixes look good to me.

[dblaikie]

In D88973#2317712, @arsenm wrote:

In D88973#2317661, @dblaikie wrote:

The goal is to recommit the old patch, with these changes in addition in the same commit? (or can this be committed first, then the other patch recommitted on top of this safely?)

I was going to squash this into the un-revert, but I could also split enough of the Attribute utility functions to pre-commit this

Ah, all good either way - just trying to understand better. Carry on :)

[arsenm]

Merged into 0a7cd99a702595ccf73c957be0127af9f25fb9a2

llvm/include/llvm/IR/Attributes.h

public:
   static Attribute getWithAllocSizeArgs(LLVMContext &Context,
                                         unsigned ElemSizeArg,
                                         const Optional<unsigned> &NumElemsArg);
   static Attribute getWithByValType(LLVMContext &Context, Type *Ty);
   static Attribute getWithStructRetType(LLVMContext &Context, Type *Ty);
   static Attribute getWithByRefType(LLVMContext &Context, Type *Ty);
   static Attribute getWithPreallocatedType(LLVMContext &Context, Type *Ty);
 
+  /// For a typed attribute, return the equivalent attribute with the type
+  /// changed to \p ReplacementTy.
+  Attribute getWithNewType(LLVMContext &Context, Type *ReplacementTy) {
+    assert(isTypeAttribute() && "this requires a typed attribute");
+    return get(Context, getKindAsEnum(), ReplacementTy);
+  }
+
   static Attribute::AttrKind getAttrKindFromName(StringRef AttrName);
 
   static StringRef getNameFromAttrKind(Attribute::AttrKind AttrKind);
 
   /// Return true if and only if the attribute has an Argument.
   static bool doesAttrKindHaveArgument(Attribute::AttrKind AttrKind);
 
   /// Return true if the provided string matches the IR name of an attribute.
public:
 
   /// Remove all attributes at the specified arg index from this
   /// attribute list. Returns a new list because attribute lists are immutable.
   LLVM_NODISCARD AttributeList removeParamAttributes(LLVMContext &C,
                                                      unsigned ArgNo) const {
     return removeAttributes(C, ArgNo + FirstArgIndex);
   }
 
+  /// Replace the type contained by attribute \p AttrKind at index \p ArgNo wih
+  /// \p ReplacementTy, preserving all other attributes.
+  LLVM_NODISCARD AttributeList replaceAttributeType(LLVMContext &C,
+                                                    unsigned ArgNo,
+                                                    Attribute::AttrKind Kind,
+                                                    Type *ReplacementTy) const {
+    Attribute Attr = getAttribute(ArgNo, Kind);
+    auto Attrs = removeAttribute(C, ArgNo, Kind);
+    return Attrs.addAttribute(C, ArgNo, Attr.getWithNewType(C, ReplacementTy));
+  }
+
   /// \brief Add the dereferenceable attribute to the attribute set at the given
   /// index. Returns a new list because attribute lists are immutable.
   LLVM_NODISCARD AttributeList addDereferenceableAttr(LLVMContext &C,
                                                       unsigned Index,
                                                       uint64_t Bytes) const;
 
   /// \brief Add the dereferenceable attribute to the attribute set at the given
   /// arg index. Returns a new list because attribute lists are immutable.

llvm/lib/Bitcode/Writer/ValueEnumerator.cpp

void ValueEnumerator::incorporateFunction(const Function &F) {
   // LocalAsMetadata.
   incorporateFunctionMetadata(F);
 
   // Adding function arguments to the value table.
   for (const auto &I : F.args()) {
     EnumerateValue(&I);
     if (I.hasAttribute(Attribute::ByVal))
       EnumerateType(I.getParamByValType());
+    else if (I.hasAttribute(Attribute::StructRet))
+      EnumerateType(I.getParamStructRetType());
   }
   FirstFuncConstantID = Values.size();
 
   // Add all function-level constants to the value table.
   for (const BasicBlock &BB : F) {
     for (const Instruction &I : BB) {
       for (const Use &OI : I.operands()) {
         if ((isa<Constant>(OI) && !isa<GlobalValue>(OI)) || isa<InlineAsm>(OI))

llvm/lib/IR/Core.cpp

LLVMAttributeRef LLVMCreateEnumAttribute(LLVMContextRef C, unsigned KindID,
   auto AttrKind = (Attribute::AttrKind)KindID;
 
   if (AttrKind == Attribute::AttrKind::ByVal) {
     // After r362128, byval attributes need to have a type attribute. Provide a
     // NULL one until a proper API is added for this.
     return wrap(Attribute::getWithByValType(Ctx, NULL));
   }
 
+  if (AttrKind == Attribute::AttrKind::StructRet) {
+    // Same as byval.
+    return wrap(Attribute::getWithStructRetType(Ctx, NULL));
+  }
+
   return wrap(Attribute::get(Ctx, AttrKind, Val));
 }
 
 unsigned LLVMGetEnumAttributeKind(LLVMAttributeRef A) {
   return unwrap(A).getKindAsEnum();
 }
 
 uint64_t LLVMGetEnumAttributeValue(LLVMAttributeRef A) {

llvm/lib/Linker/IRMover.cpp

GlobalVariable *NewDGV =
                          SGVar->getAddressSpace());
   NewDGV->setAlignment(MaybeAlign(SGVar->getAlignment()));
   NewDGV->copyAttributesFrom(SGVar);
   return NewDGV;
 }
 
 AttributeList IRLinker::mapAttributeTypes(LLVMContext &C, AttributeList Attrs) {
   for (unsigned i = 0; i < Attrs.getNumAttrSets(); ++i) {
-    if (Attrs.hasAttribute(i, Attribute::ByVal)) {
-      Type *Ty = Attrs.getAttribute(i, Attribute::ByVal).getValueAsType();
-      if (!Ty)
-        continue;
-
-      Attrs = Attrs.removeAttribute(C, i, Attribute::ByVal);
-      Attrs = Attrs.addAttribute(
-          C, i, Attribute::getWithByValType(C, TypeMap.get(Ty)));
+    for (Attribute::AttrKind TypedAttr :
+         {Attribute::ByVal, Attribute::StructRet}) {
+      if (Attrs.hasAttribute(i, TypedAttr)) {
+        if (Type *Ty = Attrs.getAttribute(i, TypedAttr).getValueAsType()) {
+          Attrs = Attrs.replaceAttributeType(C, i, TypedAttr, TypeMap.get(Ty));
+          break;
+        }
+      }
     }
   }
   return Attrs;
 }
 
 /// Link the function in the source module into the destination module if
 /// needed, setting up mapping information.
 Function *IRLinker::copyFunctionProto(const Function *SF) {

llvm/lib/Transforms/Utils/ValueMapper.cpp

if (auto *CB = dyn_cast<CallBase>(I)) {
     for (Type *Ty : FTy->params())
       Tys.push_back(TypeMapper->remapType(Ty));
     CB->mutateFunctionType(FunctionType::get(
         TypeMapper->remapType(I->getType()), Tys, FTy->isVarArg()));
 
     LLVMContext &C = CB->getContext();
     AttributeList Attrs = CB->getAttributes();
     for (unsigned i = 0; i < Attrs.getNumAttrSets(); ++i) {
-      if (Attrs.hasAttribute(i, Attribute::ByVal)) {
-        Type *Ty = Attrs.getAttribute(i, Attribute::ByVal).getValueAsType();
-        if (!Ty)
-          continue;
-
-        Attrs = Attrs.removeAttribute(C, i, Attribute::ByVal);
-        Attrs = Attrs.addAttribute(
-            C, i, Attribute::getWithByValType(C, TypeMapper->remapType(Ty)));
+      for (Attribute::AttrKind TypedAttr :
+           {Attribute::ByVal, Attribute::StructRet}) {
+        if (Type *Ty = Attrs.getAttribute(i, TypedAttr).getValueAsType()) {
+          Attrs = Attrs.replaceAttributeType(C, i, TypedAttr,
+                                             TypeMapper->remapType(Ty));
+          break;
+        }
       }
     }
     CB->setAttributes(Attrs);
     return;
   }
   if (auto *AI = dyn_cast<AllocaInst>(I))
     AI->setAllocatedType(TypeMapper->remapType(AI->getAllocatedType()));
   if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {

llvm/test/Linker/Inputs/sret-type-input.ll

+%a = type { i64 }
+%struct = type { i32, i8 }
+
+define void @g(%a* sret(%a)) {
+  ret void
+}
+
+declare void @baz(%struct* sret(%struct))
+
+define void @foo(%struct* sret(%struct) %a) {
+  call void @baz(%struct* sret(%struct) %a)
+  ret void
+}

llvm/test/Linker/sret-types.ll

+; RUN: llvm-link %s %p/Inputs/sret-type-input.ll -S | FileCheck %s
+
+%a = type { i64 }
+%struct = type { i32, i8 }
+
+; CHECK-LABEL: define void @f(%a* sret(%a) %0)
+define void @f(%a* sret(%a)) {
+  ret void
+}
+
+; CHECK-LABEL: define void @bar(
+; CHECK: call void @foo(%struct* sret(%struct) %ptr)
+define void @bar() {
+  %ptr = alloca %struct
+  call void @foo(%struct* sret(%struct) %ptr)
+  ret void
+}
+
+; CHECK-LABEL: define void @g(%a* sret(%a) %0)
+
+; CHECK-LABEL: define void @foo(%struct* sret(%struct) %a)
+; CHECK-NEXT:   call void @baz(%struct* sret(%struct) %a)
+declare void @foo(%struct* sret(%struct) %a)
+
+; CHECK: declare void @baz(%struct* sret(%struct))