[FuncSpec] Promote stack values before specialization

[mtsamis]

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:

• Since by default the specializer is ran once, without an initial call to this function worthwhile specialization opportunities may be missed. This is especially true for Fortran programs where function arguments are passed by reference.
• This also results in more consistent specialization behaviour when the specializer is ran more than once. Currently there are cases where functions will be specialized depending on if an unrelated specializable function is added.

[chill]

In https://reviews.llvm.org/D145819 this is done for each iteration, IIRC so the same transformation can be applied to the new specialised instances.

[mtsamis]

Some additional context for this change:

This is part of improving function specialization in exchange2, as initiated by https://reviews.llvm.org/D106426.
Without this change getting the critical digits_2 function to specialize requires force-specialization=true which is undesired.

In general, this helps uncover more specialization opportunities in Fortran code and results in more consistent specialization behaviour when funcspec-max-iters > 1.

[mtsamis]

In D152556#4409234, @chill wrote:

In https://reviews.llvm.org/D145819 this is done for each iteration, IIRC so the same transformation can be applied to the new specialised instances.

Thanks for the reference.

I see this moves promoteConstantStackValues() to the beggining of run which means that one the last iteration promoteConstantStackValues is not called.
This can result in some extra unpromoted / unwanted instructions after the specializer, which can be seen when funcspec-max-iters is not very large.
That's why I opted for calling it both initially and then after each iteration.

Would it be worthwhile to suggest merging this with the code on https://reviews.llvm.org/D145819?

[labrinea]

I am planning changes to D145819. It makes sense to run promoteConstantStackValues per function and not across the entire module, at least until specialization on literal constants gets enabled. The reason is that there can't be new opportunities in consecutive iterations of run() unless we are specializing a call which uses a return value from a specialization created in the previous iteration, or a call which resides in the body of a specialization. It seems to me that promoteConstantStackValues was specifically written for recursive functions, so perhaprs we could limit it there like:

@@ -476,17 +464,37 @@ bool FunctionSpecializer::run() {
     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 (Metrics.isRecursive)
+      promoteConstantStackValues(&F);

-    if (!findSpecializations(&F, SpecCost, AllSpecs, SM)) {
+    if (!findSpecializations(&F, Metrics.NumInsts, AllSpecs, SM)) {
       LLVM_DEBUG(
           dbgs() << "FnSpecialization: No possible specializations found for "
                  << F.getName() << "\n");

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

FunctionSpecializer(
       std::function<const TargetLibraryInfo &(Function &)> GetTLI,
       std::function<TargetTransformInfo &(Function &)> GetTTI,
       std::function<AssumptionCache &(Function &)> GetAC)
       : Solver(Solver), M(M), FAM(FAM), GetBFI(GetBFI), GetTLI(GetTLI),
         GetTTI(GetTTI), GetAC(GetAC) {}
 
   ~FunctionSpecializer();
 
+  void init();
   bool run();
 
   InstCostVisitor getInstCostVisitorFor(Function *F) {
     auto &BFI = GetBFI(*F);
     auto &TTI = GetTTI(*F);
     return InstCostVisitor(M.getDataLayout(), BFI, TTI, Solver);
   }
 

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

LLVM_DEBUG(
     if (NumSpecsCreated > 0)
       dbgs() << "FnSpecialization: Created " << NumSpecsCreated
              << " specializations in module " << M.getName() << "\n");
   // Eliminate dead code.
   removeDeadFunctions();
   cleanUpSSA();
 }
 
+/// Should be called once before any calls to run.
+void FunctionSpecializer::init() {
+  // Promote constant stack values once before entering the run loop.
+  // Apart from generally increasing specialization opportunities, this is
+  // critical for Fortran programs where function arguments are passed by
+  // reference.
+  promoteConstantStackValues();
+}
+
 /// Attempt to specialize functions in the module to enable constant
 /// propagation across function boundaries.
 ///
 /// \returns true if at least one function is specialized.
 bool FunctionSpecializer::run() {
   // Find possible specializations for each function.
   SpecMap SM;
   SmallVector<Spec, 32> AllSpecs;

llvm/lib/Transforms/IPO/SCCP.cpp

for (GlobalVariable &G : M.globals()) {
     if (canTrackGlobalVariableInterprocedurally(&G))
       Solver.trackValueOfGlobalVariable(&G);
   }
 
   // Solve for constants.
   Solver.solveWhileResolvedUndefsIn(M);
 
   if (IsFuncSpecEnabled) {
+    Specializer.init();
     unsigned Iters = 0;
     while (Iters++ < FuncSpecMaxIters && Specializer.run());
   }
 
   // Iterate over all of the instructions in the module, replacing them with
   // constants if we have found them to be of constant values.
   bool MadeChanges = false;
   for (Function &F : M) {

llvm/test/Transforms/FunctionSpecialization/function-specialization-recursive5.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=4 -S < %s | FileCheck %s --check-prefix=ITERS4
+
+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.3 = internal constant i32 1
+; ITERS1:  @funcspec.arg.4 = internal constant i32 5
+
+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
+;
+; 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
+;
+  %low.arg = alloca i32, align 4
+  %high.arg = alloca i32, align 4
+  store i32 0, ptr %low.arg, align 4
+  store i32 6, ptr %high.arg, align 4
+  call void @recursiveFunc(ptr nonnull %low.arg, i32 1, ptr nonnull %high.arg)
+  ret i32 0
+}
+
+declare dso_local void @print_val(i32, i32)