LLZKInlineFreeFunctionsPass.cpp

Go to the documentation of this file.

//===-- LLZKInlineFreeFunctionsPass.cpp -------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "llzk/Analysis/CallGraphAnalyses.h"
#include "llzk/Analysis/SymbolUseGraph.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
 
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Transforms/InliningUtils.h>
 
#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/SCCIterator.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/Support/Debug.h>
 
namespace llzk {
#define GEN_PASS_DEF_INLINEFREEFUNCTIONSPASS
#include "llzk/Transforms/LLZKTransformationPasses.h.inc"
} // namespace llzk
 
#define DEBUG_TYPE "llzk-inline-free-functions"
 
using namespace mlir;
using namespace llzk;
using namespace llzk::function;
 
namespace {

static bool isInlinableFreeFunction(FuncDefOp func, ModuleOp root) {
  return func->getParentOp() == root;
}

static bool isInlinableCallSite(CallOp call, ModuleOp root) {
  return call->getParentOfType<FuncDefOp>() && call->getParentOfType<ModuleOp>() == root;
}

static llvm::DenseSet<Operation *> collectRecursiveFunctions(const llzk::CallGraph &cg) {
  llvm::DenseSet<Operation *> recursive;
  for (auto scc = llvm::scc_begin(&cg); !scc.isAtEnd(); ++scc) {
    if (scc->size() == 1 && !scc.hasCycle()) {
      continue;
    }
    for (const llzk::CallGraphNode *node : *scc) {
      if (!node->isExternal()) {
        recursive.insert(node->getCalledFunction().getOperation());
      }
    }
  }
  return recursive;
}

static FuncDefOp resolveFreeCallee(
    CallOp call, ModuleOp root, SymbolTableCollection &tables,
    const llvm::DenseSet<Operation *> &skippedCallees
) {
  auto tgtRes = call.getCalleeTarget(tables);
  if (failed(tgtRes)) {
    return nullptr;
  }
  FuncDefOp callee = tgtRes->get();
  if (!isInlinableFreeFunction(callee, root) || callee.isExternal() ||
      skippedCallees.contains(callee)) {
    return nullptr;
  }
  return callee;
}

struct FreeFunctionCall {
  CallOp call;
  FuncDefOp callee;
};

static SmallVector<FreeFunctionCall> collectFreeFunctionCalls(
    ModuleOp mod, SymbolTableCollection &tables, const llvm::DenseSet<Operation *> &skippedCallees
) {
  SmallVector<FreeFunctionCall> calls;
  mod.walk([&](CallOp call) {
    if (!isInlinableCallSite(call, mod)) {
      return;
    }
    if (FuncDefOp callee = resolveFreeCallee(call, mod, tables, skippedCallees)) {
      calls.push_back({call, callee});
    }
  });
  return calls;
}

static SmallVector<FuncDefOp> collectUnusedHelpers(ModuleOp mod) {
  SymbolUseGraph useGraph(mod.getOperation());
  SmallVector<FuncDefOp> unusedFunctions;
  for (FuncDefOp func : mod.getOps<FuncDefOp>()) {
    if (func.isExternal()) {
      continue;
    }
    const SymbolUseGraphNode *node = useGraph.lookupNode(func);
    if (!node || !node->hasPredecessor()) {
      unusedFunctions.push_back(func);
    }
  }
  return unusedFunctions;
}
 
class PassImpl : public llzk::impl::InlineFreeFunctionsPassBase<PassImpl> {
  using Base = InlineFreeFunctionsPassBase<PassImpl>;
  using Base::Base;
 
  void runOnOperation() override {
    ModuleOp mod = getOperation();
    SymbolTableCollection tables;
    InlinerInterface inliner(&getContext());
    // Seeded with recursive functions; grows with callees that cannot be
    // inlined.
    llvm::DenseSet<Operation *> skippedCallees =
        collectRecursiveFunctions(getAnalysis<CallGraphAnalysis>().getCallGraph());
 
    inlineCalls(mod, tables, inliner, skippedCallees);
    removeUnusedFunctions(mod);
  }

  void inlineCalls(
      ModuleOp mod, SymbolTableCollection &tables, InlinerInterface &inliner,
      llvm::DenseSet<Operation *> &skippedCallees
  ) {
    SmallVector<FreeFunctionCall> callsToInline =
        collectFreeFunctionCalls(mod, tables, skippedCallees);
    while (!callsToInline.empty()) {
      LLVM_DEBUG({
        llvm::dbgs() << "[" DEBUG_TYPE "] round found " << callsToInline.size()
                     << " free-function call site(s) to inline\n";
      });
 
      for (auto [call, callee] : callsToInline) {
        if (skippedCallees.contains(callee)) {
          continue;
        }
        if (failed(inlineCall(inliner, call, callee, callee.getCallableRegion(), true))) {
          call.emitWarning("failed to inline free function call; skipping this callee");
          skippedCallees.insert(callee);
          continue;
        }
        call.erase();
      }
 
      callsToInline = collectFreeFunctionCalls(mod, tables, skippedCallees);
    }
  }

  void removeUnusedFunctions(ModuleOp mod) {
    SmallVector<FuncDefOp> toErase = collectUnusedHelpers(mod);
    while (!toErase.empty()) {
      for (FuncDefOp func : toErase) {
        func.erase();
      }
      toErase = collectUnusedHelpers(mod);
    }
  }
};
} // namespace