[IPSCCP] do not break musttail invariant (PR36485)

[indutny]

Summary

Do not replace results of musttail calls with a constant if the
call itself can't be removed.

Do not zap returns of musttail callees, if the call site can't be
removed and replaced with a constant.

Do not zap returns of musttail-calling blocks, this breaks
invariant too.

[indutny]

Move test file into a proper directory.

[indutny]

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

[indutny]

Removed FileCheck from the test. It's enough if it doesn't crash, there's nothing to validate.

[indutny]

I think I've "under-implemented" the fix here. First of all, it doesn't mark callees overdefined, which leads to zapped return value. Second, it basically kills whole optimization because of musttail call site.

What would be the proper implementation in this case? I've few ideas, but I could easily underestimate the scope of the problem.

[fhahn]

Thanks for the patch! Unless I miss something, we might be able to continue using the more precise lattice value while solving and only handle it when doing the actual replacement. Eli, it would be great if you could have a quick look and see if that makes sense.

Removed FileCheck from the test. It's enough if it doesn't crash, there's nothing to validate.

Could you add some checks again, checking the musttail calls are preserved? Otherwise, we could do lots of things to "not crash".

lib/Transforms/Scalar/SCCP.cpp

1256

This is not the place where the actual replacement happens. I think we could keep the lattice value for the musttail call here and continue using it for solving. The replacement happens in tryToReplaceWithConstant. We could check for musttail calls there and do the replacement there.

This way, in case we can determine the return value of a function containing musttail calls, we still can use the return value at the callsite while solving.

Also, in case the musttail callee has no side effects, the call will be removed (D38856), and we could replace the result with the constant.

[indutny]

I've addressed your feedback, and covered more erroneous cases in my patch now.

Thank you!

[efriedma]

LGTM

[indutny]

Thank you!

[efriedma]

Oh, forgot to ask; do you have commit access? If not, I can commit for you.

[fhahn]

Thanks for updating the patch and thanks Eli for having a look. LGTM too

[indutny]

I don't have commit access yet. Would greatly appreciate if you could land it for me. Thank you!

[davide]

LGTM too. Sorry I'm a little behind with reviews because life took over.

[indutny]

Thank you, everyone!

lib/Transforms/Scalar/SCCP.cpp

class SCCPSolver : public InstVisitor<SCCPSolver> {
   /// TrackedMultipleRetVals - Same as TrackedRetVals, but used for functions
   /// that return multiple values.
   DenseMap<std::pair<Function *, unsigned>, LatticeVal> TrackedMultipleRetVals;
 
   /// MRVFunctionsTracked - Each function in TrackedMultipleRetVals is
   /// represented here for efficient lookup.
   SmallPtrSet<Function *, 16> MRVFunctionsTracked;
 
+  /// MustTailFunctions - Each function here is a callee of non-removable
+  /// musttail call site.
+  SmallPtrSet<Function *, 16> MustTailCallees;
+
   /// TrackingIncomingArguments - This is the set of functions for whose
   /// arguments we make optimistic assumptions about and try to prove as
   /// constants.
   SmallPtrSet<Function *, 16> TrackingIncomingArguments;
 
   /// The reason for two worklists is that overdefined is the lowest state
   /// on the lattice, and moving things to overdefined as fast as possible
   /// makes SCCP converge much faster.
if (auto *STy = dyn_cast<StructType>(F->getReturnType())) {
       MRVFunctionsTracked.insert(F);
       for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i)
         TrackedMultipleRetVals.insert(std::make_pair(std::make_pair(F, i),
                                                      LatticeVal()));
     } else
       TrackedRetVals.insert(std::make_pair(F, LatticeVal()));
   }
 
+  /// AddMustTailCallee - If the SCCP solver finds that this function is called
+  /// from non-removable musttail call site.
+  void AddMustTailCallee(Function *F) {
+    MustTailCallees.insert(F);
+  }
+
+  /// Returns true if the given function is called from non-removable musttail
+  /// call site.
+  bool isMustTailCallee(Function *F) {
+    return MustTailCallees.count(F);
+  }
+
   void AddArgumentTrackedFunction(Function *F) {
     TrackingIncomingArguments.insert(F);
   }
 
   /// Returns true if the given function is in the solver's set of
   /// argument-tracked functions.
   bool isArgumentTrackedFunction(Function *F) {
     return TrackingIncomingArguments.count(F);
public:
   }
 
   /// getMRVFunctionsTracked - Get the set of functions which return multiple
   /// values tracked by the pass.
   const SmallPtrSet<Function *, 16> getMRVFunctionsTracked() {
     return MRVFunctionsTracked;
   }
 
+  /// getMustTailCallees - Get the set of functions which are called
+  /// from non-removable musttail call sites.
+  const SmallPtrSet<Function *, 16> getMustTailCallees() {
+    return MustTailCallees;
+  }
+
   /// markOverdefined - Mark the specified value overdefined.  This
   /// works with both scalars and structs.
   void markOverdefined(Value *V) {
     if (auto *STy = dyn_cast<StructType>(V->getType()))
       for (unsigned i = 0, e = STy->getNumElements(); i != e; ++i)
         markOverdefined(getStructValueState(V, i), V);
     else
       markOverdefined(ValueState[V], V);
if (auto *STy = dyn_cast<StructType>(F->getReturnType())) {
     if (TFRVI == TrackedRetVals.end())
       goto CallOverdefined;  // Not tracking this callee.
 
     // If so, propagate the return value of the callee into this call result.
     mergeInValue(I, TFRVI->second);
   }
 }
 
 void SCCPSolver::Solve() {
   // Process the work lists until they are empty!
   while (!BBWorkList.empty() || !InstWorkList.empty() ||
          !OverdefinedInstWorkList.empty()) {
     // Process the overdefined instruction's work list first, which drives other
     // things to overdefined more quickly.
     while (!OverdefinedInstWorkList.empty()) {
       Value *I = OverdefinedInstWorkList.pop_back_val();
 
if (IV.isConstantRange()) {
             ConstantInt::get(V->getType(), IV.asConstantInteger().getValue());
       else
         return false;
     } else
       Const =
           IV.isConstant() ? IV.getConstant() : UndefValue::get(V->getType());
   }
   assert(Const && "Constant is nullptr here!");
+
+  // Replacing `musttail` instructions with constant breaks `musttail` invariant
+  // unless the call itself can be removed
+  CallInst *CI = dyn_cast<CallInst>(V);
+  if (CI && CI->isMustTailCall() && !CI->isSafeToRemove()) {
+    CallSite CS(CI);
+    Function *F = CS.getCalledFunction();
+
+    // Don't zap returns of the callee
+    if (F)
+      Solver.AddMustTailCallee(F);
+
+    DEBUG(dbgs() << "  Can\'t treat the result of musttail call : " << *CI
+                 << " as a constant\n");
+    return false;
+  }
+
   DEBUG(dbgs() << "  Constant: " << *Const << " = " << *V << '\n');
 
   // Replaces all of the uses of a variable with uses of the constant.
   V->replaceAllUsesWith(Const);
   return true;
 }
 
 // runSCCP() - Run the Sparse Conditional Constant Propagation algorithm,

 
 static void findReturnsToZap(Function &F,
                              SmallVector<ReturnInst *, 8> &ReturnsToZap,
                              SCCPSolver &Solver) {
   // We can only do this if we know that nothing else can call the function.
   if (!Solver.isArgumentTrackedFunction(&F))
     return;
 
-  for (BasicBlock &BB : F)
+  // There is a non-removable musttail call site of this function. Zapping
+  // returns is not allowed.
+  if (Solver.isMustTailCallee(&F)) {
+    DEBUG(dbgs() << "Can't zap returns of the function : " << F.getName()
+                 << " due to present musttail call of it\n");
+    return;
+  }
+
+  for (BasicBlock &BB : F) {
+    if (CallInst *CI = BB.getTerminatingMustTailCall()) {
+      DEBUG(dbgs() << "Can't zap return of the block due to present "
+                   << "musttail call : " << *CI << "\n");
+      return;
+    }
+
     if (auto *RI = dyn_cast<ReturnInst>(BB.getTerminator()))
       if (!isa<UndefValue>(RI->getOperand(0)))
         ReturnsToZap.push_back(RI);
   }
+}
 
 static bool runIPSCCP(Module &M, const DataLayout &DL,
                       const TargetLibraryInfo *TLI) {
   SCCPSolver Solver(DL, TLI);
 
   // Loop over all functions, marking arguments to those with their addresses
   // taken or that are external as overdefined.
   for (Function &F : M) {

test/Transforms/IPConstantProp/musttail-call.ll

+; RUN: opt < %s -ipsccp -S | FileCheck %s
+; PR36485
+; musttail call result can\'t be replaced with a constant, unless the call
+; can be removed
+
+declare i32 @external()
+
+define i8* @start(i8 %v) {
+  %c1 = icmp eq i8 %v, 0
+  br i1 %c1, label %true, label %false
+true:
+  ; CHECK: %ca = musttail call i8* @side_effects(i8 %v)
+  ; CHECK: ret i8* %ca
+  %ca = musttail call i8* @side_effects(i8 %v)
+  ret i8* %ca
+false:
+  %c2 = icmp eq i8 %v, 1
+  br i1 %c2, label %c2_true, label %c2_false
+c2_true:
+  %ca1 = musttail call i8* @no_side_effects(i8 %v)
+  ; CHECK: ret i8* null
+  ret i8* %ca1
+c2_false:
+  ; CHECK: %ca2 = musttail call i8* @dont_zap_me(i8 %v)
+  ; CHECK: ret i8* %ca2
+  %ca2 = musttail call i8* @dont_zap_me(i8 %v)
+  ret i8* %ca2
+}
+
+define internal i8* @side_effects(i8 %v) {
+  %i1 = call i32 @external()
+
+  ; since this goes back to `start` the SCPP should be see that the return value
+  ; is always `null`.
+  ; The call can't be removed due to `external` call above, though.
+
+  ; CHECK: %ca = musttail call i8* @start(i8 %v)
+  %ca = musttail call i8* @start(i8 %v)
+
+  ; Thus the result must be returned anyway
+  ; CHECK: ret i8* %ca
+  ret i8* %ca
+}
+
+define internal i8* @no_side_effects(i8 %v) readonly nounwind {
+  ; The call to this function is removed, so the return value must be zapped
+  ; CHECK: ret i8* undef
+  ret i8* null
+}
+
+define internal i8* @dont_zap_me(i8 %v) {
+  %i1 = call i32 @external()
+
+  ; The call to this function cannot be removed due to side effects. Thus the
+  ; return value should stay as it is, and should not be zapped.
+  ; CHECK: ret i8* null
+  ret i8* null
+}