[FuncSpec] Promote stack values before specialization

[labrinea]

Summary

After each iteration of the function specializer, constant stack values are promoted to constant globals in order to enable recursive function specialization. This should also be done once before running the specializer. Enables specialization of _QMbrute_forcePdigits_2 from SPEC2017:548.exchange2_r

[labrinea]

Alternative to D152556

[ChuanqiXu]

LGTM. Thanks.

llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h

 private:
   Constant *getPromotableAlloca(AllocaInst *Alloca, CallInst *Call);
 
   /// A constant stack value is an AllocaInst that has a single constant
   /// value stored to it. Return this constant if such an alloca stack value
   /// is a function argument.
   Constant *getConstantStackValue(CallInst *Call, Value *Val);
 
-  /// Iterate over the argument tracked functions see if there
-  /// are any new constant values for the call instruction via
-  /// stack variables.
-  void promoteConstantStackValues();
+  /// See if there are any new constant values for the callers of \p F via
+  /// stack variables and promote them to global variables.
+  void promoteConstantStackValues(Function *F);
 
   /// Clean up fully specialized functions.
   void removeDeadFunctions();
 
   /// Remove any ssa_copy intrinsics that may have been introduced.
   void cleanUpSSA();
 
-  // Compute the code metrics for function \p F.
-  CodeMetrics &analyzeFunction(Function *F);
-
   /// @brief  Find potential specialization opportunities.
   /// @param F Function to specialize
   /// @param SpecCost Cost of specializing a function. Final score is benefit
   /// minus this cost.
   /// @param AllSpecs A vector to add potential specializations to.
   /// @param SM  A map for a function's specialisation range
   /// @return True, if any potential specializations were found
   bool findSpecializations(Function *F, Cost SpecCost,
                            SmallVectorImpl<Spec> &AllSpecs, SpecMap &SM);
 
   bool isCandidateFunction(Function *F);
 
   /// @brief Create a specialization of \p F and prime the SCCPSolver
   /// @param F Function to specialize
   /// @param S Which specialization to create
   /// @return The new, cloned function
   Function *createSpecialization(Function *F, const SpecSig &S);
 
-  /// Compute and return the cost of specializing function \p F.
-  Cost getSpecializationCost(Function *F);
-
   /// Determine if it is possible to specialise the function for constant values
   /// of the formal parameter \p A.
   bool isArgumentInteresting(Argument *A);
 
   /// Check if the value \p V  (an actual argument) is a constant or can only
   /// have a constant value. Return that constant.
   Constant *getCandidateConstant(Value *V);
 

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

 //
 //     @funcspec.arg = internal constant i32 2
 //
 //     define internal void @someFunc(i32* arg1) {
 //       call void @otherFunc(i32* nonnull @funcspec.arg)
 //       ret void
 //     }
 //
-void FunctionSpecializer::promoteConstantStackValues() {
-  // Iterate over the argument tracked functions see if there
-  // are any new constant values for the call instruction via
-  // stack variables.
-  for (Function &F : M) {
-    if (!Solver.isArgumentTrackedFunction(&F))
-      continue;
-
-    for (auto *User : F.users()) {
+// See if there are any new constant values for the callers of \p F via
+// stack variables and promote them to global variables.
+void FunctionSpecializer::promoteConstantStackValues(Function *F) {
+  for (User *U : F->users()) {
 
-      auto *Call = dyn_cast<CallInst>(User);
-      if (!Call)
-        continue;
+    auto *Call = dyn_cast<CallInst>(U);
 
-      if (!Solver.isBlockExecutable(Call->getParent()))
-        continue;
 
-      bool Changed = false;
-      for (const Use &U : Call->args()) {
-        unsigned Idx = Call->getArgOperandNo(&U);
-        Value *ArgOp = Call->getArgOperand(Idx);
-        Type *ArgOpType = ArgOp->getType();
 
-        if (!Call->onlyReadsMemory(Idx) || !ArgOpType->isPointerTy())
-          continue;
 
-        auto *ConstVal = getConstantStackValue(Call, ArgOp);
-        if (!ConstVal)
-          continue;
 
-        Value *GV = new GlobalVariable(M, ConstVal->getType(), true,
-                                       GlobalValue::InternalLinkage, ConstVal,
-                                       "funcspec.arg");
-        if (ArgOpType != ConstVal->getType())
-          GV = ConstantExpr::getBitCast(cast<Constant>(GV), ArgOpType);
 
-        Call->setArgOperand(Idx, GV);
-        Changed = true;
-      }
-
-      // Add the changed CallInst to Solver Worklist
-      if (Changed)
-        Solver.visitCall(*Call);
     }
   }
 }
 
 // ssa_copy intrinsics are introduced by the SCCP solver. These intrinsics
 // interfere with the promoteConstantStackValues() optimization.
 static void removeSSACopy(Function &F) {
   for (BasicBlock &BB : F) {
bool FunctionSpecializer::run() {
   // Find possible specializations for each function.
   SpecMap SM;
   SmallVector<Spec, 32> AllSpecs;
   unsigned NumCandidates = 0;
   for (Function &F : M) {
     if (!isCandidateFunction(&F))
       continue;
 
-    Cost SpecCost = getSpecializationCost(&F);
-    if (!SpecCost.isValid()) {
-      LLVM_DEBUG(dbgs() << "FnSpecialization: Invalid specialization cost for "
-                        << F.getName() << "\n");
-      continue;
+    auto [It, Inserted] = FunctionMetrics.try_emplace(&F);
+    CodeMetrics &Metrics = It->second;
+    //Analyze the function.
+    if (Inserted) {
+      SmallPtrSet<const Value *, 32> EphValues;
+      CodeMetrics::collectEphemeralValues(&F, &GetAC(F), EphValues);
+      for (BasicBlock &BB : F)
+        Metrics.analyzeBasicBlock(&BB, GetTTI(F), EphValues);
     }
 
+    // If the code metrics reveal that we shouldn't duplicate the function,
+    // or if the code size implies that this function is easy to get inlined,
+    // then we shouldn't specialize it.
+    if (Metrics.notDuplicatable || !Metrics.NumInsts.isValid() ||
+        (!ForceSpecialization && !F.hasFnAttribute(Attribute::NoInline) &&
+         Metrics.NumInsts < MinFunctionSize))
+      continue;
+
+    // TODO: For now only consider recursive functions when running multiple
+    // times. This should change if specialization on literal constants gets
+    // enabled.
+    if (!Inserted && !Metrics.isRecursive && !SpecializeLiteralConstant)
+      continue;
+
     LLVM_DEBUG(dbgs() << "FnSpecialization: Specialization cost for "
-                      << F.getName() << " is " << SpecCost << "\n");
+                      << F.getName() << " is " << Metrics.NumInsts << "\n");
 
-    if (!findSpecializations(&F, SpecCost, AllSpecs, SM)) {
+    if (Inserted && Metrics.isRecursive)
+      promoteConstantStackValues(&F);
+
+    if (!findSpecializations(&F, Metrics.NumInsts, AllSpecs, SM)) {
       LLVM_DEBUG(
           dbgs() << "FnSpecialization: No possible specializations found for "
                  << F.getName() << "\n");
       continue;
     }
 
     ++NumCandidates;
   }
for (User *U : F->users()) {
         Solver.resetLatticeValueFor(CS);
       }
     }
   }
 
   // Rerun the solver to notify the users of the modified callsites.
   Solver.solveWhileResolvedUndefs();
 
-  promoteConstantStackValues();
+  for (Function *F : OriginalFuncs)
+    if (FunctionMetrics[F].isRecursive)
+      promoteConstantStackValues(F);
+
   return true;
 }
 
 void FunctionSpecializer::removeDeadFunctions() {
   for (Function *F : FullySpecialized) {
     LLVM_DEBUG(dbgs() << "FnSpecialization: Removing dead function "
                       << F->getName() << "\n");
     if (FAM)
       FAM->clear(*F, F->getName());
     F->eraseFromParent();
   }
   FullySpecialized.clear();
 }
 
-// Compute the code metrics for function \p F.
-CodeMetrics &FunctionSpecializer::analyzeFunction(Function *F) {
-  auto I = FunctionMetrics.insert({F, CodeMetrics()});
-  CodeMetrics &Metrics = I.first->second;
-  if (I.second) {
-    // The code metrics were not cached.
-    SmallPtrSet<const Value *, 32> EphValues;
-    CodeMetrics::collectEphemeralValues(F, &(GetAC)(*F), EphValues);
-    for (BasicBlock &BB : *F)
-      Metrics.analyzeBasicBlock(&BB, (GetTTI)(*F), EphValues);
-  }
-  return Metrics;
-}
-
 /// Clone the function \p F and remove the ssa_copy intrinsics added by
 /// the SCCPSolver in the cloned version.
 static Function *cloneCandidateFunction(Function *F) {
   ValueToValueMapTy Mappings;
   Function *Clone = CloneFunction(F, Mappings);
   removeSSACopy(*Clone);
   return Clone;
 }
Function *FunctionSpecializer::createSpecialization(Function *F,
 
   // Mark all the specialized functions
   Specializations.insert(Clone);
   ++NumSpecsCreated;
 
   return Clone;
 }
 
-/// Compute and return the cost of specializing function \p F.
-Cost FunctionSpecializer::getSpecializationCost(Function *F) {
-  CodeMetrics &Metrics = analyzeFunction(F);
-  // If the code metrics reveal that we shouldn't duplicate the function, we
-  // shouldn't specialize it. Set the specialization cost to Invalid.
-  // Or if the lines of codes implies that this function is easy to get
-  // inlined so that we shouldn't specialize it.
-  if (Metrics.notDuplicatable || !Metrics.NumInsts.isValid() ||
-      (!ForceSpecialization && !F->hasFnAttribute(Attribute::NoInline) &&
-       Metrics.NumInsts < MinFunctionSize))
-    return InstructionCost::getInvalid();
-
-  // Otherwise, set the specialization cost to be the cost of all the
-  // instructions in the function.
-  return Metrics.NumInsts;
-}
-
 /// Compute a bonus for replacing argument \p A with constant \p C.
 Cost FunctionSpecializer::getSpecializationBonus(Argument *A, Constant *C,
                                                  InstCostVisitor &Visitor) {
   LLVM_DEBUG(dbgs() << "FnSpecialization: Analysing bonus for constant: "
                     << C->getNameOrAsOperand() << "\n");
 
   Cost TotalCost = 0;
   for (auto *U : A->users())

llvm/test/Transforms/FunctionSpecialization/function-specialization-recursive.ll

-; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=2 -S < %s | FileCheck %s --check-prefix=ITERS2
-; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=3 -S < %s | FileCheck %s --check-prefix=ITERS3
-; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=4 -S < %s | FileCheck %s --check-prefix=ITERS4
-
-@low = internal constant i32 0, align 4
-@high = internal constant i32 6, align 4
-
-define internal void @recursiveFunc(ptr nocapture readonly %lo, i32 %step, ptr nocapture readonly %hi) {
-  %lo.temp = alloca i32, align 4
-  %hi.temp = alloca i32, align 4
-  %lo.load = load i32, ptr %lo, align 4
-  %hi.load = load i32, ptr %hi, align 4
-  %cmp = icmp ne i32 %lo.load, %hi.load
-  br i1 %cmp, label %block6, label %ret.block
-
-block6:
-  call void @print_val(i32 %lo.load, i32 %hi.load)
-  %add = add nsw i32 %lo.load, %step
-  %sub = sub nsw i32 %hi.load, %step
-  store i32 %add, ptr %lo.temp, align 4
-  store i32 %sub, ptr %hi.temp, align 4
-  call void @recursiveFunc(ptr nonnull %lo.temp, i32 %step, ptr nonnull %hi.temp)
-  br label %ret.block
-
-ret.block:
-  ret void
-}
-
-; ITERS2:  @funcspec.arg.4 = internal constant i32 2
-; ITERS2:  @funcspec.arg.5 = internal constant i32 4
-
-; ITERS3:  @funcspec.arg.7 = internal constant i32 3
-; ITERS3:  @funcspec.arg.8 = internal constant i32 3
-
-define i32 @main() {
-; ITERS2-LABEL: @main(
-; ITERS2-NEXT:    call void @print_val(i32 0, i32 6)
-; ITERS2-NEXT:    call void @print_val(i32 1, i32 5)
-; ITERS2-NEXT:    call void @recursiveFunc(ptr nonnull @funcspec.arg.4, i32 1, ptr nonnull @funcspec.arg.5)
-; ITERS2-NEXT:    ret i32 0
-;
-; ITERS3-LABEL: @main(
-; ITERS3-NEXT:    call void @print_val(i32 0, i32 6)
-; ITERS3-NEXT:    call void @print_val(i32 1, i32 5)
-; ITERS3-NEXT:    call void @print_val(i32 2, i32 4)
-; ITERS3-NEXT:    call void @recursiveFunc(ptr nonnull @funcspec.arg.7, i32 1, ptr nonnull @funcspec.arg.8)
-; ITERS3-NEXT:    ret i32 0
-;
-; ITERS4-LABEL: @main(
-; ITERS4-NEXT:    call void @print_val(i32 0, i32 6)
-; ITERS4-NEXT:    call void @print_val(i32 1, i32 5)
-; ITERS4-NEXT:    call void @print_val(i32 2, i32 4)
-; ITERS4-NEXT:    ret i32 0
-;
-  call void @recursiveFunc(ptr nonnull @low, i32 1, ptr nonnull @high)
-  ret i32 0
-}
-
-declare dso_local void @print_val(i32, i32)

llvm/test/Transforms/FunctionSpecialization/promoteContantStackValues.ll

+; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=1 -S < %s | FileCheck %s --check-prefix=ITERS1
+; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=2 -S < %s | FileCheck %s --check-prefix=ITERS2
+; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=3 -S < %s | FileCheck %s --check-prefix=ITERS3
+; RUN: opt -passes="ipsccp<func-spec>,inline,instcombine" -force-specialization -funcspec-max-iters=4 -S < %s | FileCheck %s --check-prefix=ITERS4
+
+@low = internal constant i32 0, align 4
+@high = internal constant i32 6, align 4
+
+define internal void @recursiveFunc(ptr nocapture readonly %lo, i32 %step, ptr nocapture readonly %hi) {
+  %lo.temp = alloca i32, align 4
+  %hi.temp = alloca i32, align 4
+  %lo.load = load i32, ptr %lo, align 4
+  %hi.load = load i32, ptr %hi, align 4
+  %cmp = icmp ne i32 %lo.load, %hi.load
+  br i1 %cmp, label %block6, label %ret.block
+
+block6:
+  call void @print_val(i32 %lo.load, i32 %hi.load)
+  %add = add nsw i32 %lo.load, %step
+  %sub = sub nsw i32 %hi.load, %step
+  store i32 %add, ptr %lo.temp, align 4
+  store i32 %sub, ptr %hi.temp, align 4
+  call void @recursiveFunc(ptr nonnull %lo.temp, i32 %step, ptr nonnull %hi.temp)
+  br label %ret.block
+
+ret.block:
+  ret void
+}
+
+; ITERS1:  @funcspec.arg = internal constant i32 0
+; ITERS1:  @funcspec.arg.1 = internal constant i32 6
+; ITERS1:  @funcspec.arg.3 = internal constant i32 1
+; ITERS1:  @funcspec.arg.4 = internal constant i32 5
+
+; ITERS2:  @funcspec.arg = internal constant i32 0
+; ITERS2:  @funcspec.arg.1 = internal constant i32 6
+; ITERS2:  @funcspec.arg.3 = internal constant i32 1
+; ITERS2:  @funcspec.arg.4 = internal constant i32 5
+; ITERS2:  @funcspec.arg.6 = internal constant i32 2
+; ITERS2:  @funcspec.arg.7 = internal constant i32 4
+
+; ITERS3:  @funcspec.arg = internal constant i32 0
+; ITERS3:  @funcspec.arg.1 = internal constant i32 6
+; ITERS3:  @funcspec.arg.3 = internal constant i32 1
+; ITERS3:  @funcspec.arg.4 = internal constant i32 5
+; ITERS3:  @funcspec.arg.6 = internal constant i32 2
+; ITERS3:  @funcspec.arg.7 = internal constant i32 4
+; ITERS3:  @funcspec.arg.9 = internal constant i32 3
+; ITERS3:  @funcspec.arg.10 = internal constant i32 3
+
+; ITERS4:  @funcspec.arg = internal constant i32 0
+; ITERS4:  @funcspec.arg.1 = internal constant i32 6
+; ITERS4:  @funcspec.arg.3 = internal constant i32 1
+; ITERS4:  @funcspec.arg.4 = internal constant i32 5
+; ITERS4:  @funcspec.arg.6 = internal constant i32 2
+; ITERS4:  @funcspec.arg.7 = internal constant i32 4
+; ITERS4:  @funcspec.arg.9 = internal constant i32 3
+; ITERS4:  @funcspec.arg.10 = internal constant i32 3
+
+define i32 @main() {
+; ITERS1-LABEL: @main(
+; ITERS1-NEXT:    call void @print_val(i32 0, i32 6)
+; ITERS1-NEXT:    call void @recursiveFunc(ptr nonnull @funcspec.arg.3, i32 1, ptr nonnull @funcspec.arg.4)
+; ITERS1-NEXT:    ret i32 0
+;
+; ITERS2-LABEL: @main(
+; ITERS2-NEXT:    call void @print_val(i32 0, i32 6)
+; ITERS2-NEXT:    call void @print_val(i32 1, i32 5)
+; ITERS2-NEXT:    call void @recursiveFunc(ptr nonnull @funcspec.arg.6, i32 1, ptr nonnull @funcspec.arg.7)
+; ITERS2-NEXT:    ret i32 0
+;
+; ITERS3-LABEL: @main(
+; ITERS3-NEXT:    call void @print_val(i32 0, i32 6)
+; ITERS3-NEXT:    call void @print_val(i32 1, i32 5)
+; ITERS3-NEXT:    call void @print_val(i32 2, i32 4)
+; ITERS3-NEXT:    call void @recursiveFunc(ptr nonnull @funcspec.arg.9, i32 1, ptr nonnull @funcspec.arg.10)
+; ITERS3-NEXT:    ret i32 0
+;
+; ITERS4-LABEL: @main(
+; ITERS4-NEXT:    call void @print_val(i32 0, i32 6)
+; ITERS4-NEXT:    call void @print_val(i32 1, i32 5)
+; ITERS4-NEXT:    call void @print_val(i32 2, i32 4)
+; ITERS4-NEXT:    ret i32 0
+;
+  %lo.temp = alloca i32, align 4
+  %hi.temp = alloca i32, align 4
+  %lo.load = load i32, ptr @low, align 4
+  %hi.load = load i32, ptr @high, align 4
+  store i32 %lo.load, ptr %lo.temp, align 4
+  store i32 %hi.load, ptr %hi.temp, align 4
+  call void @recursiveFunc(ptr nonnull %lo.temp, i32 1, ptr nonnull %hi.temp)
+  ret i32 0
+}
+
+declare dso_local void @print_val(i32, i32)