[mlir][inliner] Refactor MLIR inliner pass and utils. (#84059)

This is just code refactoring done as a preparation for adding
MLIR inliner cost model hook(s).
Related discussion: https://discourse.llvm.org/t/inliner-cost-model/2992

The logic of SCC-based MLIR inliner is separated into the Inliner
implementation. The MLIR inliner pass becomes, well, just a pass
that invokes the SCC-based MLIR inliner.

mlir/include/mlir/Transforms/Inliner.h

@@ -0,0 +1,120 @@
+//===- Inliner.h - Inliner pass utilities -----------------------*- C++ -*-===//
+//
+// 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 header file declares utility structures for the inliner pass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_INLINER_H
+#define MLIR_TRANSFORMS_INLINER_H
+
+#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Interfaces/CallInterfaces.h"
+#include "mlir/Pass/AnalysisManager.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace mlir {
+class OpPassManager;
+class Operation;
+
+class InlinerConfig {
+public:
+  using DefaultPipelineTy = std::function<void(OpPassManager &)>;
+  using OpPipelinesTy = llvm::StringMap<OpPassManager>;
+
+  InlinerConfig() = default;
+  InlinerConfig(DefaultPipelineTy defaultPipeline,
+                unsigned maxInliningIterations)
+      : defaultPipeline(std::move(defaultPipeline)),
+        maxInliningIterations(maxInliningIterations) {}
+
+  const DefaultPipelineTy &getDefaultPipeline() const {
+    return defaultPipeline;
+  }
+  const OpPipelinesTy &getOpPipelines() const { return opPipelines; }
+  unsigned getMaxInliningIterations() const { return maxInliningIterations; }
+  void setDefaultPipeline(DefaultPipelineTy pipeline) {
+    defaultPipeline = std::move(pipeline);
+  }
+  void setOpPipelines(OpPipelinesTy pipelines) {
+    opPipelines = std::move(pipelines);
+  }
+  void setMaxInliningIterations(unsigned max) { maxInliningIterations = max; }
+
+private:
+  /// An optional function that constructs an optimization pipeline for
+  /// a given operation. This optimization pipeline is applied
+  /// only to those callable operations that do not have dedicated
+  /// optimization pipeline in opPipelines (based on the operation name).
+  DefaultPipelineTy defaultPipeline;
+  /// A map of operation names to pass pipelines to use when optimizing
+  /// callable operations of these types. This provides a specialized pipeline
+  /// instead of the one produced by defaultPipeline.
+  OpPipelinesTy opPipelines;
+  /// For SCC-based inlining algorithms, specifies maximum number of iterations
+  /// when inlining within an SCC.
+  unsigned maxInliningIterations{0};
+};
+
+/// This is an implementation of the inliner
+/// that operates bottom up over the Strongly Connected Components(SCCs)
+/// of the CallGraph. This enables a more incremental propagation
+/// of inlining decisions from the leafs to the roots of the callgraph.
+class Inliner {
+public:
+  using RunPipelineHelperTy = std::function<LogicalResult(
+      Pass &pass, OpPassManager &pipeline, Operation *op)>;
+
+  Inliner(Operation *op, CallGraph &cg, Pass &pass, AnalysisManager am,
+          RunPipelineHelperTy runPipelineHelper, const InlinerConfig &config)
+      : op(op), cg(cg), pass(pass), am(am),
+        runPipelineHelper(std::move(runPipelineHelper)), config(config) {}
+  Inliner(Inliner &) = delete;
+  void operator=(const Inliner &) = delete;
+
+  /// Perform inlining on a OpTrait::SymbolTable operation.
+  LogicalResult doInlining();
+
+  /// This struct represents a resolved call to a given callgraph node. Given
+  /// that the call does not actually contain a direct reference to the
+  /// Region(CallGraphNode) that it is dispatching to, we need to resolve them
+  /// explicitly.
+  struct ResolvedCall {
+    ResolvedCall(CallOpInterface call, CallGraphNode *sourceNode,
+                 CallGraphNode *targetNode)
+        : call(call), sourceNode(sourceNode), targetNode(targetNode) {}
+    CallOpInterface call;
+    CallGraphNode *sourceNode, *targetNode;
+  };
+
+protected:
+  /// An OpTrait::SymbolTable operation to run the inlining on.
+  Operation *op;
+  /// A CallGraph analysis for the given operation.
+  CallGraph &cg;
+  /// A reference to the pass using this inliner.
+  Pass &pass;
+  /// Analysis manager for the given operation instance.
+  AnalysisManager am;
+  /// A callback for running a nested pass pipeline on the operation
+  /// contained within the main operation.
+  const RunPipelineHelperTy runPipelineHelper;
+  /// The inliner configuration parameters.
+  const InlinerConfig &config;
+
+private:
+  /// Forward declaration of the class providing the actual implementation.
+  class Impl;
+
+public:
+};
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_INLINER_H

mlir/include/mlir/Transforms/Passes.td

@@ -269,7 +269,9 @@ def Inliner : Pass<"inline"> {
   let constructor = "mlir::createInlinerPass()";
   let options = [
     Option<"defaultPipelineStr", "default-pipeline", "std::string",
-           /*default=*/"\"canonicalize\"", "The default optimizer pipeline used for callables">,
+           /*default=*/"\"canonicalize\"",
+           "The optimizer pipeline used for callables that do not have "
+           "a dedicated optimizer pipeline in opPipelineList">,
     ListOption<"opPipelineList", "op-pipelines", "OpPassManager",
                "Callable operation specific optimizer pipelines (in the form "
                "of `dialect.op(pipeline)`)">,

mlir/lib/Transforms/CMakeLists.txt

@@ -5,7 +5,7 @@ add_mlir_library(MLIRTransforms
   ControlFlowSink.cpp
   CSE.cpp
   GenerateRuntimeVerification.cpp
-  Inliner.cpp
+  InlinerPass.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   Mem2Reg.cpp

mlir/lib/Transforms/InlinerPass.cpp

@@ -0,0 +1,155 @@
+//===- InlinerPass.cpp - Pass to inline function calls --------------------===//
+//
+// 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 file implements a basic inlining algorithm that operates bottom up over
+// the Strongly Connect Components(SCCs) of the CallGraph. This enables a more
+// incremental propagation of inlining decisions from the leafs to the roots of
+// the callgraph.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Analysis/CallGraph.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Inliner.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_INLINER
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+/// This function implements the inliner optimization pipeline.
+static void defaultInlinerOptPipeline(OpPassManager &pm) {
+  pm.addPass(createCanonicalizerPass());
+}
+
+//===----------------------------------------------------------------------===//
+// InlinerPass
+//===----------------------------------------------------------------------===//
+
+namespace {
+class InlinerPass : public impl::InlinerBase<InlinerPass> {
+public:
+  InlinerPass();
+  InlinerPass(const InlinerPass &) = default;
+  InlinerPass(std::function<void(OpPassManager &)> defaultPipeline);
+  InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
+              llvm::StringMap<OpPassManager> opPipelines);
+  void runOnOperation() override;
+
+  /// A callback provided to the inliner driver to execute
+  /// the specified pass pipeline on the given operation
+  /// within the context of the current inliner pass,
+  /// which is passed as the first argument.
+  /// runPipeline API is protected within the Pass class,
+  /// so this helper is required to call it from the foreign
+  /// inliner driver.
+  static LogicalResult runPipelineHelper(Pass &pass, OpPassManager &pipeline,
+                                         Operation *op) {
+    return mlir::cast<InlinerPass>(pass).runPipeline(pipeline, op);
+  }
+
+private:
+  /// Attempt to initialize the options of this pass from the given string.
+  /// Derived classes may override this method to hook into the point at which
+  /// options are initialized, but should generally always invoke this base
+  /// class variant.
+  LogicalResult initializeOptions(StringRef options) override;
+
+  /// Inliner configuration parameters created from the pass options.
+  InlinerConfig config;
+};
+} // namespace
+
+InlinerPass::InlinerPass() : InlinerPass(defaultInlinerOptPipeline) {}
+
+InlinerPass::InlinerPass(
+    std::function<void(OpPassManager &)> defaultPipelineArg)
+    : InlinerPass(std::move(defaultPipelineArg),
+                  llvm::StringMap<OpPassManager>{}) {}
+
+InlinerPass::InlinerPass(std::function<void(OpPassManager &)> defaultPipeline,
+                         llvm::StringMap<OpPassManager> opPipelines)
+    : config(std::move(defaultPipeline), maxInliningIterations) {
+  if (opPipelines.empty())
+    return;
+
+  // Update the option for the op specific optimization pipelines.
+  for (auto &it : opPipelines)
+    opPipelineList.addValue(it.second);
+  config.setOpPipelines(std::move(opPipelines));
+}
+
+void InlinerPass::runOnOperation() {
+  CallGraph &cg = getAnalysis<CallGraph>();
+
+  // The inliner should only be run on operations that define a symbol table,
+  // as the callgraph will need to resolve references.
+  Operation *op = getOperation();
+  if (!op->hasTrait<OpTrait::SymbolTable>()) {
+    op->emitOpError() << " was scheduled to run under the inliner, but does "
+                         "not define a symbol table";
+    return signalPassFailure();
+  }
+
+  // Get an instance of the inliner.
+  Inliner inliner(op, cg, *this, getAnalysisManager(), runPipelineHelper,
+                  config);
+
+  // Run the inlining.
+  if (failed(inliner.doInlining()))
+    signalPassFailure();
+  return;
+}
+
+LogicalResult InlinerPass::initializeOptions(StringRef options) {
+  if (failed(Pass::initializeOptions(options)))
+    return failure();
+
+  // Initialize the pipeline builder for operations without the dedicated
+  // optimization pipeline in opPipelineList to use the option string.
+  // TODO: Use a generic pass manager for the pre-inline pipeline, and remove
+  // this.
+  if (!defaultPipelineStr.empty()) {
+    std::string defaultPipelineCopy = defaultPipelineStr;
+    config.setDefaultPipeline([=](OpPassManager &pm) {
+      (void)parsePassPipeline(defaultPipelineCopy, pm);
+    });
+  } else if (defaultPipelineStr.getNumOccurrences()) {
+    config.setDefaultPipeline(nullptr);
+  }
+
+  // Initialize the op specific pass pipelines.
+  llvm::StringMap<OpPassManager> pipelines;
+  for (OpPassManager pipeline : opPipelineList)
+    if (!pipeline.empty())
+      pipelines.try_emplace(pipeline.getOpAnchorName(), pipeline);
+  config.setOpPipelines(std::move(pipelines));
+
+  config.setMaxInliningIterations(maxInliningIterations);
+
+  return success();
+}
+
+std::unique_ptr<Pass> mlir::createInlinerPass() {
+  return std::make_unique<InlinerPass>();
+}
+std::unique_ptr<Pass>
+mlir::createInlinerPass(llvm::StringMap<OpPassManager> opPipelines) {
+  return std::make_unique<InlinerPass>(defaultInlinerOptPipeline,
+                                       std::move(opPipelines));
+}
+std::unique_ptr<Pass> mlir::createInlinerPass(
+    llvm::StringMap<OpPassManager> opPipelines,
+    std::function<void(OpPassManager &)> defaultPipelineBuilder) {
+  return std::make_unique<InlinerPass>(std::move(defaultPipelineBuilder),
+                                       std::move(opPipelines));
+}

mlir/lib/Transforms/Utils/CMakeLists.txt

@@ -5,6 +5,7 @@ add_mlir_library(MLIRTransformUtils
   DialectConversion.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
+  Inliner.cpp
   InliningUtils.cpp
   LoopInvariantCodeMotionUtils.cpp
   OneToNTypeConversion.cpp