[Attributor][FIX] Prevent alignment breakage wrt. must-tail calls

[jdoerfert]

Summary

If we have a must-tail call the callee and caller need to have matching
ABIs. Part of that is alignment which we might modify when we deduce
alignment of arguments of either. Since we would need to keep them in
sync, which is not as simple, we simply avoid deducing alignment for
arguments of the must-tail caller or callee.

[lebedev.ri]

How do we know this is only about alignment?

[jdoerfert]

In D76673#1938580, @lebedev.ri wrote:

How do we know this is only about alignment?

Good point. It is about:
Attribute::StructRet, Attribute::ByVal, Attribute::InAlloca, Attribute::InReg, Attribute::Returned, Attribute::SwiftSelf, Attribute::SwiftError, Attribute::Alignment

We derive Returned and Alignment so I need to add the prevention to returned as well. I'll also move it to the manifest as we can use the derived value internally just fine.

[jdoerfert]

Derive but do not manifest alignment and returned for (call site) arguments

[lebedev.ri]

Thanks, this looks about reasonable to me, although i'm not familiar with this code to fully review it.

llvm/lib/Transforms/IPO/Attributor.cpp

1364

We are in AAReturnedValues

[jdoerfert]

llvm/lib/Transforms/IPO/Attributor.cpp

1364

You know, copy&paste ;)

[uenoku]

Good point. It is about:
Attribute::StructRet, Attribute::ByVal, Attribute::InAlloca, Attribute::InReg, Attribute::Returned, Attribute::SwiftSelf, Attribute::SwiftError, Attribute::Alignment

Seems reasonable but I'm not sure why we need to care about returned.

[jdoerfert]

In D76673#1941116, @uenoku wrote:

Good point. It is about:
Attribute::StructRet, Attribute::ByVal, Attribute::InAlloca, Attribute::InReg, Attribute::Returned, Attribute::SwiftSelf, Attribute::SwiftError, Attribute::Alignment

Seems reasonable but I'm not sure why we need to care about returned.

It's in the verifier's list that is checked. @rnk can you confirm returned needs to be kept in-sync between caller and callee of a must-tail call?

[rnk]

lgtm

I think the current verifier check is overly conservative. I think the align argument attribute is only an ABI-relevant attribute when it is used with byval. In your test case, it seems to refer to the alignment of the pointee of a pointer argument, not the alignment of the argument itself. In that case, it is not ABI relevant, and we could tolerate the mismatch.

However, I just don't think it's worth worrying about this detail in attributor. Mainly, it points out that align is overloaded. =(

[jdoerfert]

In D76673#1945260, @rnk wrote:

lgtm

I think the current verifier check is overly conservative. I think the align argument attribute is only an ABI-relevant attribute when it is used with byval. In your test case, it seems to refer to the alignment of the pointee of a pointer argument, not the alignment of the argument itself. In that case, it is not ABI relevant, and we could tolerate the mismatch.

However, I just don't think it's worth worrying about this detail in attributor. Mainly, it points out that align is overloaded. =(

I don't think in my test it is the pointee of the pointer argument. The load from %p has an alignment of 32, so %p is assumed to be align 32 at this location. Then the argument of musttail_callee_1 is assumed to be align 32 as well but we do not manifest it.
What am I getting wrong?

Also, the verifier check lists returned as well. Is it correct that returned has ABI implications?

[rnk]

In D76673#1945275, @jdoerfert wrote:

In D76673#1945260, @rnk wrote:

lgtm

I think the current verifier check is overly conservative. I think the align argument attribute is only an ABI-relevant attribute when it is used with byval. In your test case, it seems to refer to the alignment of the pointee of a pointer argument, not the alignment of the argument itself. In that case, it is not ABI relevant, and we could tolerate the mismatch.

However, I just don't think it's worth worrying about this detail in attributor. Mainly, it points out that align is overloaded. =(

I don't think in my test it is the pointee of the pointer argument. The load from %p has an alignment of 32, so %p is assumed to be align 32 at this location. Then the argument of musttail_callee_1 is assumed to be align 32 as well but we do not manifest it.
What am I getting wrong?

Sorry, I'm not being clear. Suppose an argument is passed in memory. What is the alignment of that memory? When the attributes byval(<ty>) align N are used, it indicates the alignment of the argument memory.

From the perspective of LLVM IR, byval arguments are represented with a pointer SSA value, so the align attribute means almost the same thing: it means some number of low bits of this pointer are zero.

But from the perspective of call lowering, the align attribute can either be a statement about the bits of the argument value, or a statement about how to lay out arguments in memory on the stack, depending on the presence of other attributes. I tend to see things from this perspective, not the perspective of an IR optimizer. :)

Also, the verifier check lists returned as well. Is it correct that returned has ABI implications?

I think it's just an optimization hint. I believe the IR has to actually return the value in question, so it would be semantics-preserving to discard all returned attributes. I could be wrong, though.

[jdoerfert]

In D76673#1946852, @rnk wrote:

In D76673#1945275, @jdoerfert wrote:

In D76673#1945260, @rnk wrote:

lgtm

I think the current verifier check is overly conservative. I think the align argument attribute is only an ABI-relevant attribute when it is used with byval. In your test case, it seems to refer to the alignment of the pointee of a pointer argument, not the alignment of the argument itself. In that case, it is not ABI relevant, and we could tolerate the mismatch.

However, I just don't think it's worth worrying about this detail in attributor. Mainly, it points out that align is overloaded. =(

I don't think in my test it is the pointee of the pointer argument. The load from %p has an alignment of 32, so %p is assumed to be align 32 at this location. Then the argument of musttail_callee_1 is assumed to be align 32 as well but we do not manifest it.
What am I getting wrong?

Sorry, I'm not being clear. Suppose an argument is passed in memory. What is the alignment of that memory? When the attributes byval(<ty>) align N are used, it indicates the alignment of the argument memory.

From the perspective of LLVM IR, byval arguments are represented with a pointer SSA value, so the align attribute means almost the same thing: it means some number of low bits of this pointer are zero.

But from the perspective of call lowering, the align attribute can either be a statement about the bits of the argument value, or a statement about how to lay out arguments in memory on the stack, depending on the presence of other attributes. I tend to see things from this perspective, not the perspective of an IR optimizer. :)

Fair, but the verifier works on IR and complains about align without byval ;)
Anyway. I'm fine with not-optimizing this for the moment to appease the verifier for now.

Also, the verifier check lists returned as well. Is it correct that returned has ABI implications?

I think it's just an optimization hint. I believe the IR has to actually return the value in question, so it would be semantics-preserving to discard all returned attributes. I could be wrong, though.

So you think we can relax the verifier check which enforces the returned attribute?

---

Long story short, do you have any concerns with this?

[rnk]

In D76673#1947072, @jdoerfert wrote:

So you think we can relax the verifier check which enforces the returned attribute?

Yes.

Long story short, do you have any concerns with this?

No, go for it, I forgot to push the button.

llvm/include/llvm/Transforms/IPO/Attributor.h

TargetLibraryInfo *getTargetLibraryInfoForFunction(const Function &F) {
     return AG.getAnalysis<TargetLibraryAnalysis>(F);
   }
 
   /// Return AliasAnalysis Result for function \p F.
   AAResults *getAAResultsForFunction(const Function &F) {
     return AG.getAnalysis<AAManager>(F);
   }
 
+  /// Return true if \p Arg is involved in a must-tail call, thus the argument
+  /// of the caller or callee.
+  bool isInvolvedInMustTailCall(const Argument &Arg) const {
+    return FunctionsCalledViaMustTail.count(Arg.getParent()) ||
+           FunctionsWithMustTailCall.count(Arg.getParent());
+  }
+
   /// Return the analysis result from a pass \p AP for function \p F.
   template <typename AP>
   typename AP::Result *getAnalysisResultForFunction(const Function &F) {
     return AG.getAnalysis<AP>(F);
   }
 
   /// Return SCC size on call graph for function \p F or 0 if unknown.
   unsigned getSccSize(const Function &F) {
private:
 
   /// A nested map that remembers all instructions in a function with a certain
   /// instruction opcode (Instruction::getOpcode()).
   FuncInstOpcodeMapTy FuncInstOpcodeMap;
 
   /// A map from functions to their instructions that may read or write memory.
   FuncRWInstsMapTy FuncRWInstsMap;
 
+  /// Functions called by a `musttail` call.
+  SmallPtrSet<Function *, 8> FunctionsCalledViaMustTail;
+
+  /// Functions containing a `musttail` call.
+  SmallPtrSet<Function *, 8> FunctionsWithMustTailCall;
+
   /// The datalayout used in the module.
   const DataLayout &DL;
 
   /// MustBeExecutedContextExplorer
   MustBeExecutedContextExplorer Explorer;
 
   /// A map with knowledge retained in `llvm.assume` instructions.
   RetainedKnowledgeMap KnowledgeMap;

llvm/lib/IR/Verifier.cpp

static bool isTypeCongruent(Type *L, Type *R) {
   PointerType *PR = dyn_cast<PointerType>(R);
   if (!PL || !PR)
     return false;
   return PL->getAddressSpace() == PR->getAddressSpace();
 }
 
 static AttrBuilder getParameterABIAttributes(int I, AttributeList Attrs) {
   static const Attribute::AttrKind ABIAttrs[] = {
-      Attribute::StructRet, Attribute::ByVal, Attribute::InAlloca,
-      Attribute::InReg, Attribute::Returned, Attribute::SwiftSelf,
-      Attribute::SwiftError};
+      Attribute::StructRet, Attribute::ByVal,     Attribute::InAlloca,
+      Attribute::InReg,     Attribute::SwiftSelf, Attribute::SwiftError};
   AttrBuilder Copy;
   for (auto AK : ABIAttrs) {
     if (Attrs.hasParamAttribute(I, AK))
       Copy.addAttribute(AK);
   }
   // `align` is ABI-affecting only in combination with `byval`.
   if (Attrs.hasParamAttribute(I, Attribute::Alignment) &&
       Attrs.hasParamAttribute(I, Attribute::ByVal))

llvm/lib/Transforms/IPO/Attributor.cpp

ChangeStatus AAReturnedValuesImpl::manifest(Attributor &A) {
   };
 
   // If the assumed unique return value is an argument, annotate it.
   if (auto *UniqueRVArg = dyn_cast<Argument>(UniqueRV.getValue())) {
     // TODO: This should be handled differently!
     this->AnchorVal = UniqueRVArg;
     this->KindOrArgNo = UniqueRVArg->getArgNo();
     Changed = IRAttribute::manifest(A);
   } else if (auto *RVC = dyn_cast<Constant>(UniqueRV.getValue())) {
     // We can replace the returned value with the unique returned constant.
     Value &AnchorValue = getAnchorValue();
     if (Function *F = dyn_cast<Function>(&AnchorValue)) {
       for (const Use &U : F->uses())
         if (CallBase *CB = dyn_cast<CallBase>(U.getUser()))
           if (CB->isCallee(&U)) {
             Constant *RVCCast =
                 CB->getType() == RVC->getType()
struct AAAlignReturned final
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override { STATS_DECLTRACK_FNRET_ATTR(aligned) }
 };
 
 /// Align attribute for function argument.
 struct AAAlignArgument final
     : AAArgumentFromCallSiteArgumentsAndMustBeExecutedContext<AAAlign,
                                                               AAAlignImpl> {
-  AAAlignArgument(const IRPosition &IRP)
-      : AAArgumentFromCallSiteArgumentsAndMustBeExecutedContext<AAAlign,
-                                                                AAAlignImpl>(
-            IRP) {}
+  using Base =
+      AAArgumentFromCallSiteArgumentsAndMustBeExecutedContext<AAAlign,
+                                                              AAAlignImpl>;
+  AAAlignArgument(const IRPosition &IRP) : Base(IRP) {}
+
+  /// See AbstractAttribute::manifest(...).
+  ChangeStatus manifest(Attributor &A) override {
+    // If the associated argument is involved in a must-tail call we give up
+    // because we would need to keep the argument alignments of caller and
+    // callee in-sync. Just does not seem worth the trouble right now.
+    if (A.getInfoCache().isInvolvedInMustTailCall(*getAssociatedArgument()))
+      return ChangeStatus::UNCHANGED;
+    return Base::manifest(A);
+  }
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override { STATS_DECLTRACK_ARG_ATTR(aligned) }
 };
 
 struct AAAlignCallSiteArgument final : AAAlignFloating {
   AAAlignCallSiteArgument(const IRPosition &IRP) : AAAlignFloating(IRP) {}
 
   /// See AbstractAttribute::manifest(...).
   ChangeStatus manifest(Attributor &A) override {
+    // If the associated argument is involved in a must-tail call we give up
+    // because we would need to keep the argument alignments of caller and
+    // callee in-sync. Just does not seem worth the trouble right now.
+    if (Argument *Arg = getAssociatedArgument())
+      if (A.getInfoCache().isInvolvedInMustTailCall(*Arg))
+        return ChangeStatus::UNCHANGED;
     ChangeStatus Changed = AAAlignImpl::manifest(A);
     MaybeAlign InheritAlign =
         getAssociatedValue().getPointerAlignment(A.getDataLayout());
     if (InheritAlign.valueOrOne() >= getAssumedAlign())
       Changed = ChangeStatus::UNCHANGED;
     return Changed;
   }
 
for (Instruction &I : instructions(&F)) {
     // Note: There are no concrete attributes now so this is initially empty.
     switch (I.getOpcode()) {
     default:
       assert((!ImmutableCallSite(&I)) && (!isa<CallBase>(&I)) &&
              "New call site/base instruction type needs to be known in the "
              "Attributor.");
       break;
     case Instruction::Call:
-      // Calls are interesting but for `llvm.assume` calls we also fill the
-      // KnowledgeMap as we find them.
+      // Calls are interesting on their own, additionally:
+      // For `llvm.assume` calls we also fill the KnowledgeMap as we find them.
+      // For `must-tail` calls we remember the caller and callee.
       if (IntrinsicInst *Assume = dyn_cast<IntrinsicInst>(&I)) {
         if (Assume->getIntrinsicID() == Intrinsic::assume)
           fillMapFromAssume(*Assume, InfoCache.KnowledgeMap);
+      } else if (cast<CallInst>(I).isMustTailCall()) {
+        InfoCache.FunctionsWithMustTailCall.insert(&F);
+        InfoCache.FunctionsCalledViaMustTail.insert(
+            cast<CallInst>(I).getCalledFunction());
       }
       LLVM_FALLTHROUGH;
-    case Instruction::Load:
-      // The alignment of a pointer is interesting for loads.
-    case Instruction::Store:
-      // The alignment of a pointer is interesting for stores.
     case Instruction::CallBr:
     case Instruction::Invoke:
     case Instruction::CleanupRet:
     case Instruction::CatchSwitch:
     case Instruction::AtomicRMW:
     case Instruction::AtomicCmpXchg:
     case Instruction::Br:
     case Instruction::Resume:
     case Instruction::Ret:
+    case Instruction::Load:
+      // The alignment of a pointer is interesting for loads.
+    case Instruction::Store:
+      // The alignment of a pointer is interesting for stores.
       IsInterestingOpcode = true;
     }
     if (IsInterestingOpcode)
       InstOpcodeMap[I.getOpcode()].push_back(&I);
     if (I.mayReadOrWriteMemory())
       ReadOrWriteInsts.push_back(&I);
   }
 

 }
 
 void Attributor::identifyDefaultAbstractAttributes(Function &F) {
   if (!VisitedFunctions.insert(&F).second)
     return;
   if (F.isDeclaration())
     return;
 
+  // In non-module runs we need to look at the call sites of a function to
+  // determine if it is part of a must-tail call edge. This will influence what
+  // attributes we can derive.
+  if (!isModulePass() && !InfoCache.FunctionsCalledViaMustTail.count(&F))
+    for (const Use &U : F.uses())
+      if (ImmutableCallSite ICS = ImmutableCallSite(U.getUser()))
+        if (ICS.isCallee(&U) && ICS.isMustTailCall())
+          InfoCache.FunctionsCalledViaMustTail.insert(&F);
+
   IRPosition FPos = IRPosition::function(F);
 
   // Check for dead BasicBlocks in every function.
   // We need dead instruction detection because we do not want to deal with
   // broken IR in which SSA rules do not apply.
   getOrCreateAAFor<AAIsDead>(FPos);
 
   // Every function might be "will-return".

llvm/test/Transforms/Attributor/align.ll

 ; ATTRIBUTOR-NEXT:    [[BC:%.*]] = bitcast i32* [[ARG]] to i8*
 ; ATTRIBUTOR-NEXT:    store i8 0, i8* [[BC]], align 2048
 ; ATTRIBUTOR-NEXT:    ret void
 ;
   %bc = bitcast i32* %arg to i8*
   store i8 0, i8* %bc
   ret void
 }
+
+; Make sure we do not annotate the callee of a must-tail call with an alignment
+; we cannot also put on the caller.
+@cnd = external global i1
+define i32 @musttail_callee_1(i32* %p) {
+  %v = load i32, i32* %p, align 32
+  ret i32 %v
+}
+define i32 @musttail_caller_1(i32* %p) {
+  %c = load i1, i1* @cnd
+  br i1 %c, label %mt, label %exit
+mt:
+  %v = musttail call i32 @musttail_callee_1(i32* %p)
+  ret i32 %v
+exit:
+  ret i32 0
+}
 ; UTC_ARGS: --disable
 
 attributes #0 = { nounwind uwtable noinline }
 attributes #1 = { uwtable noinline }
 attributes #2 = { "null-pointer-is-valid"="true" }