[FuncSpec] Make the Function Specializer part of the IPSCCP pass.

Reland 877a9f9 since D138654 (parent)
has been fixed with 9ebaf4f and with
8f1e11c.

Differential Revision: https://reviews.llvm.org/D126455

llvm/include/llvm/InitializePasses.h

@@ -147,7 +147,6 @@ void initializeFlattenCFGLegacyPassPass(PassRegistry &);
 void initializeFloat2IntLegacyPassPass(PassRegistry&);
 void initializeForceFunctionAttrsLegacyPassPass(PassRegistry&);
 void initializeFuncletLayoutPass(PassRegistry&);
-void initializeFunctionSpecializationLegacyPassPass(PassRegistry &);
 void initializeGCMachineCodeAnalysisPass(PassRegistry&);
 void initializeGCModuleInfoPass(PassRegistry&);
 void initializeGVNHoistLegacyPassPass(PassRegistry&);

llvm/include/llvm/LinkAllPasses.h

@@ -217,7 +217,6 @@ namespace {
       (void) llvm::createInjectTLIMappingsLegacyPass();
       (void) llvm::createUnifyLoopExitsPass();
       (void) llvm::createFixIrreduciblePass();
-      (void)llvm::createFunctionSpecializationPass();
       (void)llvm::createSelectOptimizePass();
 
       (void)new llvm::IntervalPartition();

llvm/include/llvm/Transforms/IPO.h

@@ -146,11 +146,6 @@ ModulePass *createDeadArgHackingPass();
 ///
 ModulePass *createIPSCCPPass();
 
-//===----------------------------------------------------------------------===//
-/// createFunctionSpecializationPass - This pass propagates constants from call
-/// sites to the specialized version of the callee function.
-ModulePass *createFunctionSpecializationPass();
-
 //===----------------------------------------------------------------------===//
 //
 /// createLoopExtractorPass - This pass extracts all natural loops from the

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

@@ -0,0 +1,174 @@
+//===- FunctionSpecialization.h - Function Specialization -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This specialises functions with constant parameters. Constant parameters
+// like function pointers and constant globals are propagated to the callee by
+// specializing the function. The main benefit of this pass at the moment is
+// that indirect calls are transformed into direct calls, which provides inline
+// opportunities that the inliner would not have been able to achieve. That's
+// why function specialisation is run before the inliner in the optimisation
+// pipeline; that is by design. Otherwise, we would only benefit from constant
+// passing, which is a valid use-case too, but hasn't been explored much in
+// terms of performance uplifts, cost-model and compile-time impact.
+//
+// Current limitations:
+// - It does not yet handle integer ranges. We do support "literal constants",
+//   but that's off by default under an option.
+// - The cost-model could be further looked into (it mainly focuses on inlining
+//   benefits),
+//
+// Ideas:
+// - With a function specialization attribute for arguments, we could have
+//   a direct way to steer function specialization, avoiding the cost-model,
+//   and thus control compile-times / code-size.
+//
+// Todos:
+// - Specializing recursive functions relies on running the transformation a
+//   number of times, which is controlled by option
+//   `func-specialization-max-iters`. Thus, increasing this value and the
+//   number of iterations, will linearly increase the number of times recursive
+//   functions get specialized, see also the discussion in
+//   https://reviews.llvm.org/D106426 for details. Perhaps there is a
+//   compile-time friendlier way to control/limit the number of specialisations
+//   for recursive functions.
+// - Don't transform the function if function specialization does not trigger;
+//   the SCCPSolver may make IR changes.
+//
+// References:
+// - 2021 LLVM Dev Mtg “Introducing function specialisation, and can we enable
+//   it by default?”, https://www.youtube.com/watch?v=zJiCjeXgV5Q
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_IPO_FUNCTIONSPECIALIZATION_H
+#define LLVM_TRANSFORMS_IPO_FUNCTIONSPECIALIZATION_H
+
+#include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Transforms/Scalar/SCCP.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/SCCPSolver.h"
+#include "llvm/Transforms/Utils/SizeOpts.h"
+
+using namespace llvm;
+
+namespace llvm {
+// Bookkeeping struct to pass data from the analysis and profitability phase
+// to the actual transform helper functions.
+struct SpecializationInfo {
+  SmallVector<ArgInfo, 8> Args; // Stores the {formal,actual} argument pairs.
+  InstructionCost Gain;         // Profitability: Gain = Bonus - Cost.
+  Function *Clone;              // The definition of the specialized function.
+};
+
+using CallSpecBinding = std::pair<CallBase *, SpecializationInfo>;
+// We are using MapVector because it guarantees deterministic iteration
+// order across executions.
+using SpecializationMap = SmallMapVector<CallBase *, SpecializationInfo, 8>;
+
+class FunctionSpecializer {
+
+  /// The IPSCCP Solver.
+  SCCPSolver &Solver;
+
+  Module &M;
+
+  /// Analysis manager, needed to invalidate analyses.
+  FunctionAnalysisManager *FAM;
+
+  /// Analyses used to help determine if a function should be specialized.
+  std::function<const TargetLibraryInfo &(Function &)> GetTLI;
+  std::function<TargetTransformInfo &(Function &)> GetTTI;
+  std::function<AssumptionCache &(Function &)> GetAC;
+
+  // The number of functions specialised, used for collecting statistics and
+  // also in the cost model.
+  unsigned NbFunctionsSpecialized = 0;
+
+  SmallPtrSet<Function *, 32> SpecializedFuncs;
+  SmallPtrSet<Function *, 32> FullySpecialized;
+  DenseMap<Function *, CodeMetrics> FunctionMetrics;
+
+public:
+  FunctionSpecializer(
+      SCCPSolver &Solver, Module &M, FunctionAnalysisManager *FAM,
+      std::function<const TargetLibraryInfo &(Function &)> GetTLI,
+      std::function<TargetTransformInfo &(Function &)> GetTTI,
+      std::function<AssumptionCache &(Function &)> GetAC)
+      : Solver(Solver), M(M), FAM(FAM), GetTLI(GetTLI), GetTTI(GetTTI),
+        GetAC(GetAC) {}
+
+  ~FunctionSpecializer() {
+    // Eliminate dead code.
+    removeDeadFunctions();
+    cleanUpSSA();
+  }
+
+  bool isClonedFunction(Function *F) { return SpecializedFuncs.count(F); }
+
+  bool run();
+
+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();
+
+  /// 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);
+
+  /// This function decides whether it's worthwhile to specialize function
+  /// \p F based on the known constant values its arguments can take on. It
+  /// only discovers potential specialization opportunities without actually
+  /// applying them.
+  ///
+  /// \returns true if any specializations have been found.
+  bool findSpecializations(Function *F, InstructionCost Cost,
+                           SmallVectorImpl<CallSpecBinding> &WorkList);
+
+  bool isCandidateFunction(Function *F);
+
+  Function *createSpecialization(Function *F, CallSpecBinding &Specialization);
+
+  /// Compute and return the cost of specializing function \p F.
+  InstructionCost getSpecializationCost(Function *F);
+
+  /// Compute a bonus for replacing argument \p A with constant \p C.
+  InstructionCost getSpecializationBonus(Argument *A, Constant *C,
+                                         const LoopInfo &LI);
+
+  /// 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);
+
+  /// Redirects callsites of function \p F to its specialized copies.
+  void updateCallSites(Function *F,
+                       SmallVectorImpl<CallSpecBinding> &Specializations);
+};
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_FUNCTIONSPECIALIZATION_H

llvm/include/llvm/Transforms/IPO/SCCP.h

@@ -32,14 +32,6 @@ class IPSCCPPass : public PassInfoMixin<IPSCCPPass> {
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
-/// Pass to perform interprocedural constant propagation by specializing
-/// functions
-class FunctionSpecializationPass
-    : public PassInfoMixin<FunctionSpecializationPass> {
-public:
-  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
-};
-
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_SCCP_H

llvm/include/llvm/Transforms/Scalar/SCCP.h

@@ -40,12 +40,6 @@ class SCCPPass : public PassInfoMixin<SCCPPass> {
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
-bool runFunctionSpecialization(
-    Module &M, FunctionAnalysisManager *FAM, const DataLayout &DL,
-    std::function<TargetLibraryInfo &(Function &)> GetTLI,
-    std::function<TargetTransformInfo &(Function &)> GetTTI,
-    std::function<AssumptionCache &(Function &)> GetAC,
-    function_ref<AnalysisResultsForFn(Function &)> GetAnalysis);
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SCCP_H

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

@@ -118,6 +118,10 @@ class SCCPSolver {
   /// should be rerun.
   bool resolvedUndefsIn(Function &F);
 
+  void solveWhileResolvedUndefsIn(Module &M);
+
+  void solveWhileResolvedUndefsIn(SmallVectorImpl<Function *> &WorkList);
+
   bool isBlockExecutable(BasicBlock *BB) const;
 
   // isEdgeFeasible - Return true if the control flow edge from the 'From' basic

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -267,10 +267,6 @@ static cl::opt<bool> EnableConstraintElimination(
     cl::desc(
         "Enable pass to eliminate conditions based on linear constraints"));
 
-static cl::opt<bool> EnableFunctionSpecialization(
-    "enable-function-specialization", cl::init(false), cl::Hidden,
-    cl::desc("Enable Function Specialization pass"));
-
 static cl::opt<AttributorRunOption> AttributorRun(
     "attributor-enable", cl::Hidden, cl::init(AttributorRunOption::NONE),
     cl::desc("Enable the attributor inter-procedural deduction pass"),
@@ -1016,10 +1012,6 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
   for (auto &C : PipelineEarlySimplificationEPCallbacks)
     C(MPM, Level);
 
-  // Specialize functions with IPSCCP.
-  if (EnableFunctionSpecialization && Level == OptimizationLevel::O3)
-    MPM.addPass(FunctionSpecializationPass());
-
   // Interprocedural constant propagation now that basic cleanup has occurred
   // and prior to optimizing globals.
   // FIXME: This position in the pipeline hasn't been carefully considered in
@@ -1634,8 +1626,6 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
     MPM.addPass(PGOIndirectCallPromotion(
         true /* InLTO */, PGOOpt && PGOOpt->Action == PGOOptions::SampleUse));
 
-    if (EnableFunctionSpecialization && Level == OptimizationLevel::O3)
-      MPM.addPass(FunctionSpecializationPass());
     // Propagate constants at call sites into the functions they call.  This
     // opens opportunities for globalopt (and inlining) by substituting function
     // pointers passed as arguments to direct uses of functions.

llvm/lib/Passes/PassRegistry.def

@@ -59,7 +59,6 @@ MODULE_PASS("elim-avail-extern", EliminateAvailableExternallyPass())
 MODULE_PASS("extract-blocks", BlockExtractorPass())
 MODULE_PASS("forceattrs", ForceFunctionAttrsPass())
 MODULE_PASS("function-import", FunctionImportPass())
-MODULE_PASS("function-specialization", FunctionSpecializationPass())
 MODULE_PASS("globaldce", GlobalDCEPass())
 MODULE_PASS("globalopt", GlobalOptPass())
 MODULE_PASS("globalsplit", GlobalSplitPass())