[FuncSpec] Replace markArgInFuncSpecialization with handleCallArguments.

[labrinea]

Summary

This will propagate constant arguments to the body of specializations that we would have otherwise missed, i.e literal constants since specializing such arguments is disabled by default.

[chill]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

843

I understand that with this change we will propagate literal constant arguments into the specialized functions. But, going forward, we intend to enable funcspec-for-literal-constant by default and then this functions here has the potential to be more efficient.

Thus I would suggest not deleting it for now.

[labrinea]

llvm/lib/Transforms/Utils/SCCPSolver.cpp

843

markArgInFuncSpecialization() does not handle struct arguments as is. I fixed that in D145374, but on the review we agreed replacing it with handleCallArguments() is cleaner. Perhaps I could rework that patch and drop this one then.

[labrinea]

Put markArgInFuncSpecialization back in.

llvm/include/llvm/Transforms/Utils/SCCPSolver.h

public:
   /// to the cloned function and \p Args contains a list of constant arguments
   /// represented as pairs of {formal,actual} values (the formal argument is
   /// associated with the original function definition). All other arguments of
   /// the specialization inherit the lattice state of their corresponding values
   /// in the original function.
   void markArgInFuncSpecialization(Function *F,
                                    const SmallVectorImpl<ArgInfo> &Args);
 
+  /// Computes the lattice value for each argument of the function called by
+  /// \p CB
+  void handleCallArguments(CallBase &CB);
+
   /// Mark all of the blocks in function \p F non-executable. Clients can used
   /// this method to erase a function from the module (e.g., if it has been
   /// completely specialized and is no longer needed).
   void markFunctionUnreachable(Function *F);
 
   void visit(Instruction *I);
   void visitCall(CallInst &I);
 

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

bool FunctionSpecializer::run() {
 
   // Create the chosen specializations.
   SmallPtrSet<Function *, 8> OriginalFuncs;
   SmallVector<Function *> Clones;
   for (unsigned I = 0; I < NSpecs; ++I) {
     Spec &S = AllSpecs[BestSpecs[I]];
     S.Clone = createSpecialization(S.F, S.Sig);
 
+    // Visiting the callsites of an argument tracked function will mark
+    // its entry basic block as executable. Since there are no known
+    // callsites at this point, we need to do this manually.
+    if (S.CallSites.empty())
+      Solver.markBlockExecutable(&S.Clone->front());
+
     // Update the known call sites to call the clone.
     for (CallBase *Call : S.CallSites) {
       LLVM_DEBUG(dbgs() << "FnSpecialization: Redirecting " << *Call
                         << " to call " << S.Clone->getName() << "\n");
       Call->setCalledFunction(S.Clone);
+      Solver.handleCallArguments(*Call);
     }
 
     Clones.push_back(S.Clone);
     OriginalFuncs.insert(S.F);
   }
 
   Solver.solveWhileResolvedUndefsIn(Clones);
 
bool FunctionSpecializer::isCandidateFunction(Function *F) {
   LLVM_DEBUG(dbgs() << "FnSpecialization: Try function: " << F->getName()
                     << "\n");
   return true;
 }
 
 Function *FunctionSpecializer::createSpecialization(Function *F, const SpecSig &S) {
   Function *Clone = cloneCandidateFunction(F);
 
-  // Initialize the lattice state of the arguments of the function clone,
-  // marking the argument on which we specialized the function constant
-  // with the given value.
-  Solver.markArgInFuncSpecialization(Clone, S.Args);
-
   Solver.addArgumentTrackedFunction(Clone);
-  Solver.markBlockExecutable(&Clone->front());
 
   // Mark all the specialized functions
   Specializations.insert(Clone);
   ++NumSpecsCreated;
 
   return Clone;
 }
 
for (const Spec &S : make_range(Begin, End)) {
 
       BestSpec = &S;
     }
 
     if (BestSpec) {
       LLVM_DEBUG(dbgs() << "FnSpecialization: Redirecting " << *CS
                         << " to call " << BestSpec->Clone->getName() << "\n");
       CS->setCalledFunction(BestSpec->Clone);
+      Solver.handleCallArguments(*CS);
       ShouldDecrementCount = true;
     }
 
     if (ShouldDecrementCount)
       --NCallsLeft;
   }
 
   // If the function has been completely specialized, the original function
   // is no longer needed. Mark it unreachable.
   if (NCallsLeft == 0) {
     Solver.markFunctionUnreachable(F);
     FullySpecialized.insert(F);
   }
 }

llvm/lib/Transforms/Utils/SCCPSolver.cpp

if (Iter != AdditionalUsers.end()) {
         if (auto *UI = dyn_cast<Instruction>(U))
           ToNotify.push_back(UI);
       for (Instruction *UI : ToNotify)
         operandChangedState(UI);
     }
   }
   void handleCallOverdefined(CallBase &CB);
   void handleCallResult(CallBase &CB);
-  void handleCallArguments(CallBase &CB);
   void handleExtractOfWithOverflow(ExtractValueInst &EVI,
                                    const WithOverflowInst *WO, unsigned Idx);
 
 private:
   friend class InstVisitor<SCCPInstVisitor>;
 
   // visit implementations - Something changed in this instruction.  Either an
   // operand made a transition, or the instruction is newly executable.  Change
public:
 
   SmallPtrSetImpl<Function *> &getArgumentTrackedFunctions() {
     return TrackingIncomingArguments;
   }
 
   void markArgInFuncSpecialization(Function *F,
                                    const SmallVectorImpl<ArgInfo> &Args);
 
+  void handleCallArguments(CallBase &CB);
+
   void markFunctionUnreachable(Function *F) {
     for (auto &BB : *F)
       BBExecutable.erase(&BB);
   }
 
   void solveWhileResolvedUndefsIn(Module &M) {
     bool ResolvedUndefs = true;
     while (ResolvedUndefs) {
if (SCCPSolver::isOverdefined(LV))
       return nullptr;
     Const = SCCPSolver::isConstant(LV) ? getConstant(LV)
                                        : UndefValue::get(V->getType());
   }
   assert(Const && "Constant is nullptr here!");
   return Const;
 }
 
 void SCCPInstVisitor::markArgInFuncSpecialization(
     Function *F, const SmallVectorImpl<ArgInfo> &Args) {
   assert(!Args.empty() && "Specialization without arguments");
   assert(F->arg_size() == Args[0].Formal->getParent()->arg_size() &&
          "Functions should have the same number of arguments");
 
   auto Iter = Args.begin();
   Argument *NewArg = F->arg_begin();
   Argument *OldArg = Args[0].Formal->getParent()->arg_begin();
SmallPtrSetImpl<Function *> &SCCPSolver::getArgumentTrackedFunctions() {
   return Visitor->getArgumentTrackedFunctions();
 }
 
 void SCCPSolver::markArgInFuncSpecialization(
     Function *F, const SmallVectorImpl<ArgInfo> &Args) {
   Visitor->markArgInFuncSpecialization(F, Args);
 }
 
+void SCCPSolver::handleCallArguments(CallBase &CB) {
+  Visitor->handleCallArguments(CB);
+}
+
 void SCCPSolver::markFunctionUnreachable(Function *F) {
   Visitor->markFunctionUnreachable(F);
 }
 
 void SCCPSolver::visit(Instruction *I) { Visitor->visit(I); }
 
 void SCCPSolver::visitCall(CallInst &I) { Visitor->visitCall(I); }

llvm/test/Transforms/FunctionSpecialization/identical-specializations.ll

merge:
   %cmp2 = call i64 @compute(i64 %ph, i64 42, ptr @plus, ptr @minus)
   ret i64 %cmp2
 }
 
 define internal i64 @compute(i64 %x, i64 %y, ptr %binop1, ptr %binop2) {
 entry:
   %op0 = call i64 %binop1(i64 %x, i64 %y)
   %op1 = call i64 %binop2(i64 %x, i64 %y)
-  %op2 = call i64 @compute(i64 %x, i64 %y, ptr %binop1, ptr @plus)
+  %op2 = call i64 @compute(i64 %x, i64 42, ptr %binop1, ptr @plus)
   %add0 = add i64 %op0, %op1
   %add1 = add i64 %add0, %op2
   %div = sdiv i64 %add1, %x
   %sub = sub i64 %div, %y
   %mul = mul i64 %sub, 2
   ret i64 %mul
 }
 
 define internal i64 @plus(i64 %x, i64 %y) {
 ; CHECK-LABEL: @plus(
 entry:
   %add = add i64 %x, %y
   ret i64 %add
 }
 
 define internal i64 @minus(i64 %x, i64 %y) {
 ; CHECK-LABEL: @minus(
 entry:
   %sub = sub i64 %x, %y
   ret i64 %sub
 }
 
 ; CHECK-LABEL: @compute.1
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CMP0:%.*]] = call i64 %binop1(i64 [[X:%.*]], i64 [[Y:%.*]])
-; CHECK-NEXT:    [[CMP1:%.*]] = call i64 @plus(i64 [[X]], i64 [[Y]])
-; CHECK-NEXT:    [[CMP2:%.*]] = call i64 @compute.1(i64 [[X]], i64 [[Y]], ptr %binop1, ptr @plus)
+; CHECK-NEXT:    [[OP0:%.*]] = call i64 %binop1(i64 [[X:%.*]], i64 42)
+; CHECK-NEXT:    [[OP1:%.*]] = call i64 @plus(i64 [[X]], i64 42)
+; CHECK-NEXT:    [[OP2:%.*]] = call i64 @compute.1(i64 [[X]], i64 42, ptr %binop1, ptr @plus)
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[OP0]], [[OP1]]
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[ADD0]], [[OP2]]
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv i64 [[ADD1]], [[X]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[DIV]], 42
+; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[SUB]], 2
+; CHECK-NEXT:    ret i64 [[MUL]]
 
 ; CHECK-LABEL: @compute.2
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CMP0:%.*]] = call i64 @plus(i64 [[X:%.*]], i64 [[Y:%.*]])
-; CHECK-NEXT:    [[CMP1:%.*]] = call i64 @minus(i64 [[X]], i64 [[Y]])
-; CHECK-NEXT:    [[CMP2:%.*]] = call i64 @compute.1(i64 [[X]], i64 [[Y]], ptr @plus, ptr @plus)
+; CHECK-NEXT:    [[OP0:%.*]] = call i64 @plus(i64 [[X:%.*]], i64 [[Y:%y]])
+; CHECK-NEXT:    [[OP1:%.*]] = call i64 @minus(i64 [[X]], i64 [[Y]])
+; CHECK-NEXT:    [[OP2:%.*]] = call i64 @compute.1(i64 [[X]], i64 42, ptr @plus, ptr @plus)
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[OP0]], [[OP1]]
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[ADD0]], [[OP2]]
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv i64 [[ADD1]], [[X]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[DIV]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[SUB]], 2
+; CHECK-NEXT:    ret i64 [[MUL]]
 
 ; CHECK-LABEL: @compute.3
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CMP0:%.*]] = call i64 @minus(i64 [[X:%.*]], i64 [[Y:%.*]])
-; CHECK-NEXT:    [[CMP1:%.*]] = call i64 @plus(i64 [[X]], i64 [[Y]])
-; CHECK-NEXT:    [[CMP2:%.*]] = call i64 @compute.3(i64 [[X]], i64 [[Y]], ptr @minus, ptr @plus)
+; CHECK-NEXT:    [[OP0:%.*]] = call i64 @minus(i64 [[X:%.*]], i64 [[Y:%y]])
+; CHECK-NEXT:    [[OP1:%.*]] = call i64 @plus(i64 [[X]], i64 [[Y]])
+; CHECK-NEXT:    [[OP2:%.*]] = call i64 @compute.3(i64 [[X]], i64 42, ptr @minus, ptr @plus)
+; CHECK-NEXT:    [[ADD0:%.*]] = add i64 [[OP0]], [[OP1]]
+; CHECK-NEXT:    [[ADD1:%.*]] = add i64 [[ADD0]], [[OP2]]
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv i64 [[ADD1]], [[X]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i64 [[DIV]], [[Y]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul i64 [[SUB]], 2
+; CHECK-NEXT:    ret i64 [[MUL]]
+

llvm/test/Transforms/FunctionSpecialization/specialize-multiple-arguments.ll

 ; THREE-NEXT:    [[TMP3:%.+]] = sdiv i64 [[TMP2]], %x
 ; THREE-NEXT:    [[TMP4:%.+]] = sub i64 [[TMP3]], %y
 ; THREE-NEXT:    [[TMP5:%.+]] = mul i64 [[TMP4]], 2
 ; THREE-NEXT:    ret i64 [[TMP5]]
 ; THREE-NEXT:  }
 ;
 ; THREE-LABEL: define internal i64 @compute.3(i64 %x, i64 %y, ptr %binop1, ptr %binop2) {
 ; THREE-NEXT:  entry:
-; THREE-NEXT:    [[TMP0:%.+]] = call i64 @minus(i64 %x, i64 %y)
-; THREE-NEXT:    [[TMP1:%.+]] = call i64 @power(i64 %x, i64 %y)
+; THREE-NEXT:    [[TMP0:%.+]] = call i64 @minus(i64 %x, i64 42)
+; THREE-NEXT:    [[TMP1:%.+]] = call i64 @power(i64 %x, i64 42)
 ; THREE-NEXT:    [[TMP2:%.+]] = add i64 [[TMP0]], [[TMP1]]
 ; THREE-NEXT:    [[TMP3:%.+]] = sdiv i64 [[TMP2]], %x
-; THREE-NEXT:    [[TMP4:%.+]] = sub i64 [[TMP3]], %y
+; THREE-NEXT:    [[TMP4:%.+]] = sub i64 [[TMP3]], 42
 ; THREE-NEXT:    [[TMP5:%.+]] = mul i64 [[TMP4]], 2
 ; THREE-NEXT:    ret i64 [[TMP5]]
 ; THREE-NEXT:  }
 ;
 define internal i64 @compute(i64 %x, i64 %y, ptr %binop1, ptr %binop2) {
 entry:
   %tmp0 = call i64 %binop1(i64 %x, i64 %y)
   %tmp1 = call i64 %binop2(i64 %x, i64 %y)