[DAE] don't remove args of musttail target/caller

[indutny]

Summary

musttail requires identical signatures of caller and callee. Removing
arguments breaks musttail semantics.

PR36441

[indutny]

A question: what is backport policy of LLVM project? I'd love to see this fix in 5.0.0 .

[indutny]

Use getTerminatingMustTailCall instead of iterating through all instructions.

[aprantl]

I'm not sure I know enough about the subject matter to give more than a superficial review.

lib/Transforms/IPO/DeadArgumentElimination.cpp

512

range-based-for?

test/Transforms/DeadArgElim/musttail-caller.ll

3

It would be better to write a comment here explaining what behavior is being tested instead of just referring to the PR>.

[indutny]

lib/Transforms/IPO/DeadArgumentElimination.cpp

512

FWIW, I've just added the braces to the existing code. It seemed to be out-of-scope change when I was working on it. Should I do it now?

[indutny]

Added explanation to the test.

[rnk]

Thanks for the fix!

test/Transforms/DeadArgElim/musttail-caller.ll

6

Do we pass undef to the call? That would be correct. It doesn't matter what we pass, so long as we match the prototype.

[indutny]

test/Transforms/DeadArgElim/musttail-caller.ll

6

I'm not sure I got your comment. Do you want me to pass undef below?

[rnk]

test/Transforms/DeadArgElim/musttail-caller.ll

6

No. Can you add CHECKs for the musttail call to show what arguments are passed here? I'm wondering if we get or miss that optimization. It is invalid to change the prototype of the foo, but it is valid to pass undef because the args are dead.

[indutny]

test/Transforms/DeadArgElim/musttail-caller.ll

6

I believe it is musttail call void @foo(i32 %a, i32 0)

[indutny]

Add extra check to the test.

[rnk]

lgtm with nit

test/Transforms/DeadArgElim/musttail-caller.ll

6

Can you add a FIXME that we should replace these with undef?

[indutny]

Add FIXME.

[indutny]

Thank you for review!

[indutny]

@rnk not sure if it matters, but I don't have commit access yet. If you could land the change - I'd appreciate it a lot. Thank you!

[indutny]

Thank you, everyone!

lib/Transforms/IPO/DeadArgumentElimination.cpp

void DeadArgumentEliminationPass::SurveyFunction(const Function &F) {
 
   using RetUses = SmallVector<UseVector, 5>;
 
   // These vectors map each return value to the uses that make it MaybeLive, so
   // we can add those to the Uses map if the return value really turns out to be
   // MaybeLive. Initialized to a list of RetCount empty lists.
   RetUses MaybeLiveRetUses(RetCount);
 
-  for (Function::const_iterator BB = F.begin(), E = F.end(); BB != E; ++BB)
-    if (const ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator()))
+  bool HasMustTailCalls = false;
+
+  for (Function::const_iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
+    if (const ReturnInst *RI = dyn_cast<ReturnInst>(BB->getTerminator())) {
       if (RI->getNumOperands() != 0 && RI->getOperand(0)->getType()
           != F.getFunctionType()->getReturnType()) {
         // We don't support old style multiple return values.
         MarkLive(F);
         return;
       }
+    }
+
+    // If we have any returns of `musttail` results - the signature can't
+    // change
+    if (BB->getTerminatingMustTailCall() != nullptr)
+      HasMustTailCalls = true;
+  }
+
+  if (HasMustTailCalls) {
+    DEBUG(dbgs() << "DeadArgumentEliminationPass - " << F.getName()
+                 << " has musttail calls\n");
+  }
 
   if (!F.hasLocalLinkage() && (!ShouldHackArguments || F.isIntrinsic())) {
     MarkLive(F);
     return;
   }
 
   DEBUG(dbgs() << "DeadArgumentEliminationPass - Inspecting callers for fn: "
                << F.getName() << "\n");
   // Keep track of the number of live retvals, so we can skip checks once all
   // of them turn out to be live.
   unsigned NumLiveRetVals = 0;
+
+  bool HasMustTailCallers = false;
+
   // Loop all uses of the function.
   for (const Use &U : F.uses()) {
     // If the function is PASSED IN as an argument, its address has been
     // taken.
     ImmutableCallSite CS(U.getUser());
     if (!CS || !CS.isCallee(&U)) {
       MarkLive(F);
       return;
     }
 
+    // The number of arguments for `musttail` call must match the number of
+    // arguments of the caller
+    if (CS.isMustTailCall())
+      HasMustTailCallers = true;
+
     // If this use is anything other than a call site, the function is alive.
     const Instruction *TheCall = CS.getInstruction();
     if (!TheCall) {   // Not a direct call site?
       MarkLive(F);
       return;
     }
 
     // If we end up here, we are looking at a direct call to our function.
for (const Use &U : TheCall->uses()) {
               MaybeLiveRetUses[i].append(MaybeLiveAggregateUses.begin(),
                                          MaybeLiveAggregateUses.end());
           }
         }
       }
     }
   }
 
+  if (HasMustTailCallers) {
+    DEBUG(dbgs() << "DeadArgumentEliminationPass - " << F.getName()
+                 << " has musttail callers\n");
+  }
+
   // Now we've inspected all callers, record the liveness of our return values.
   for (unsigned i = 0; i != RetCount; ++i)
     MarkValue(CreateRet(&F, i), RetValLiveness[i], MaybeLiveRetUses[i]);
 
   DEBUG(dbgs() << "DeadArgumentEliminationPass - Inspecting args for fn: "
                << F.getName() << "\n");
 
   // Now, check all of our arguments.
   unsigned i = 0;
   UseVector MaybeLiveArgUses;
   for (Function::const_arg_iterator AI = F.arg_begin(),
        E = F.arg_end(); AI != E; ++AI, ++i) {
     Liveness Result;
-    if (F.getFunctionType()->isVarArg()) {
+    if (F.getFunctionType()->isVarArg() || HasMustTailCallers ||
+        HasMustTailCalls) {
       // Variadic functions will already have a va_arg function expanded inside
       // them, making them potentially very sensitive to ABI changes resulting
       // from removing arguments entirely, so don't. For example AArch64 handles
       // register and stack HFAs very differently, and this is reflected in the
       // IR which has already been generated.
+      //
+      // `musttail` calls to this function restrict argument removal attempts.
+      // The signature of the caller must match the signature of the function.
+      //
+      // `musttail` calls in this function prevents us from changing its
+      // signature
       Result = Live;
     } else {
       // See what the effect of this use is (recording any uses that cause
       // MaybeLive in MaybeLiveArgUses). 
       Result = SurveyUses(&*AI, MaybeLiveArgUses);
     }
 
     // Mark the result.

test/Transforms/DeadArgElim/musttail-caller.ll

+; RUN: opt -deadargelim -S < %s | FileCheck %s
+; PR36441
+; Dead arguments should not be removed in presence of `musttail` calls.
+; The signature of caller must always match the signature of the callee.
+
+; CHECK-LABEL: define internal void @test(i32 %a, i32 %b)
+define internal void @test(i32 %a, i32 %b) {
+  musttail call void @foo(i32 %a, i32 0)
+  ret void
+}
+
+; CHECK-LABEL: define internal void @foo(i32 %a, i32 %b)
+define internal void @foo(i32 %a, i32 %b) {
+  ret void
+}