[clang] [llvm] [Transforms] Implement always_specialize attribute lowering (PR #143983)

[Adrian Vogelsgesang via cfe-commits]

================
@@ -0,0 +1,324 @@
+//===- AlwaysSpecializer.cpp - implementation of always_specialize --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Function specialisation under programmer control.
+//
+// Specifically, function parameters are marked [[always_specialize]], then call
+// sites which pass a constant argument are rewritten to call specialisations.
+//
+// The difficult parts of function specialisation are the cost model, ensuring
+// termination and specialisation to the anticipated extent.
+//
+// Cost model is under programmer control, exactly like always_inline.
+//
+// Termination follows from the implementation following a phased structure:
+// 1. Functions are identifed in the input IR
+// 2. Calls that exist in the input IR are identified
+// Those constitute the complete set of specialisations that will be created.
+//
+// This pass does the _minimum_ specialisation, in the sense that only call
+// sites in the input will lead to cloning. A specialised function will call
+// another specialised function iff there was a call site with the same
+// argument vector in the input.
+//
+// Running the identifyCalls + createClones sequence N times will behave
+// as expected, specialising recursively to that depth. This patch has N=1
+// in the first instance, with no commandline argument to override.
+// Similarly variadic functions are not yet handled.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/IPO/AlwaysSpecializer.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/InstructionSimplify.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/IPO/FunctionSpecialization.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "always-specialize"
+
+namespace {
+
+class AlwaysSpecializer : public ModulePass {
+public:
+  static char ID;
+
+  AlwaysSpecializer() : ModulePass(ID) {}
+  StringRef getPassName() const override { return "Always specializer"; }
+
+  // One constant for each argument, nullptr if that one is non-constant
+  using ArgVector = SmallVector<Constant *, 4>;
+
+  // A map from the ArgVector to the matching specialisation
+  using FunctionSpecializations = MapVector<ArgVector, Function *>;
+
+  // The four mini-passes populate and then use a map:
+  // 1. identifyFunctions writes all keys, with default initialised values.
+  // 2. identifyCalls writes all the ArgVector keys in the values of SpecList.
+  // 3. createClones writes the Function* values at the leaves.
+  // 4. replaceCalls walks the map doing the trivial rewrite.
+
+  // Conceptually a Map<Function*, Specialization> but a vector suffices.
+  using SpecListTy =
+      SmallVector<std::pair<Function *, FunctionSpecializations>, 4>;
+
+  SpecListTy identifyFunctions(Module &M);
+  bool identifyCalls(Module &M, Function *F, FunctionSpecializations &);
+  bool createClones(Module &M, Function *F, FunctionSpecializations &);
+  bool replaceCalls(Module &M, Function *F, FunctionSpecializations &);
+
+  bool runOnModule(Module &M) override {
+    bool Changed = false;
+
+    // Sets all the keys in the structure used in this invocation.
+    SpecListTy SpecList = identifyFunctions(M);
+    size_t Count = SpecList.size();
+    if (Count == 0) {
+      return false;
+    }
+
+    // Record distinct call sites as vector<Constant*> -> nullptr
+    for (auto &[F, spec] : SpecList)
+      Changed |= identifyCalls(M, F, spec);
+
+    // Create and record the clones. Note that call sites within the clones
+    // cannot trigger creating more clones so no termination risk.
+    for (auto &[F, spec] : SpecList)
+      Changed |= createClones(M, F, spec);
+
+    // Replacing calls as the final phase means no need to track
+    // partially-specialised calls and no creating further clones.
+    for (auto &[F, spec] : SpecList)
+      Changed |= replaceCalls(M, F, spec);
+
+    return Changed;
+  }
+
+  static bool isCandidateFunction(const Function &F);
+  static bool callEligible(const Function &F, const CallBase *CB,
+                           ArgVector &Out);
+  static Function *cloneCandidateFunction(Module &M, Function *F,
+                                          const ArgVector &C);
+
+  // Only a member variable to reuse the allocation. Short lived.
+  ArgVector ArgVec;
+};
+
+AlwaysSpecializer::SpecListTy AlwaysSpecializer::identifyFunctions(Module &M) {
+  SpecListTy SpecList;
+  for (Function &F : M) {
+    if (isCandidateFunction(F)) {
+      SpecList.push_back(std::make_pair(&F, FunctionSpecializations()));
+    }
+  }
+  return SpecList;
+}
+
+bool AlwaysSpecializer::identifyCalls(Module &M, Function *F,
+                                      FunctionSpecializations &Specs) {
+  bool Found = false;
+
+  for (User *U : F->users()) {
+    CallBase *CB = dyn_cast<CallBase>(U);
+    if (!CB || !callEligible(*F, CB, ArgVec)) {
+      continue;
+    }
+
+    if (!Specs.contains(ArgVec)) {
+      Found = true;
+      Specs.insert(std::make_pair(ArgVec, nullptr));
+    }
----------------
vogelsgesang wrote:

should we also emit optimization remarks?

"Function `foobar` was not specialized because its 3rd parameter (`size`) is non-constant"
"Function `foobar` was specialized with 3rd parameter `size = 15`"

https://github.com/llvm/llvm-project/pull/143983