InlineEnzymeRegions.cpp

Go to the documentation of this file.

//===- InlineEnzymeRegions.cpp - Inline/outline enzyme.autodiff  ------------ //
//
// 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 passes to inlining and outlining to convert
// between enzyme.autodiff and enzyme.autodiff_region ops.
//
//===----------------------------------------------------------------------===//
#include "Dialect/Ops.h"
#include "Interfaces/AutoDiffOpInterface.h"
#include "Passes/Passes.h"
 
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/TypeSwitch.h"
 
using namespace mlir;
 
namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_INLINEENZYMEINTOREGIONPASS
#define GEN_PASS_DEF_OUTLINEENZYMEFROMREGIONPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
namespace {
constexpr static llvm::StringLiteral kFnAttrsName = "fn_attrs";
 
static StringRef getFunctionTypeAttrName(Operation *operation) {
  if (auto iface = dyn_cast<mlir::enzyme::AutoDiffFunctionInterface>(operation))
    return iface.getFunctionTypeAttrName();
  return "";
}
 
static StringRef getArgAttrsAttrName(Operation *operation) {
  if (auto iface = dyn_cast<mlir::enzyme::AutoDiffFunctionInterface>(operation))
    return iface.getArgAttrsAttrName();
  return "";
}
 
static void serializeFunctionAttributes(Operation *fn, Operation *regionOp) {
  SmallVector<NamedAttribute> fnAttrs;
  fnAttrs.reserve(fn->getAttrDictionary().size());
  for (auto attr : fn->getAttrs()) {
    // Don't store the function type or sym_name because they may change when
    // outlining
    if (attr.getName() == getFunctionTypeAttrName(fn) ||
        attr.getName() == SymbolTable::getSymbolAttrName())
      continue;
    fnAttrs.push_back(attr);
  }
 
  regionOp->setAttr(kFnAttrsName,
                    DictionaryAttr::getWithSorted(fn->getContext(), fnAttrs));
}
 
template <typename DiffRegionOp>
static void deserializeFunctionAttributes(DiffRegionOp op,
                                          Operation *outlinedFunc,
                                          unsigned addedArgCount) {
  if (!op->template hasAttrOfType<DictionaryAttr>(kFnAttrsName))
    return;
 
  MLIRContext *ctx = op->getContext();
  SmallVector<NamedAttribute> fnAttrs;
  for (auto attr : op->template getAttrOfType<DictionaryAttr>(kFnAttrsName)) {
    // New arguments are potentially added when outlining due to references to
    // values outside the region. Insert an empty arg attr for each newly
    // added argument.
    if (attr.getName() == getArgAttrsAttrName(outlinedFunc)) {
      SmallVector<Attribute> argAttrs(
          cast<ArrayAttr>(attr.getValue())
              .template getAsRange<DictionaryAttr>());
      for (unsigned i = 0; i < addedArgCount; ++i)
        argAttrs.push_back(DictionaryAttr::getWithSorted(ctx, {}));
      fnAttrs.push_back(
          NamedAttribute(attr.getName(), ArrayAttr::get(ctx, argAttrs)));
    } else
      fnAttrs.push_back(attr);
  }
  outlinedFunc->setAttrs(fnAttrs);
}
 
struct InlineEnzymeAutoDiff : public OpRewritePattern<enzyme::AutoDiffOp> {
  using OpRewritePattern<enzyme::AutoDiffOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(enzyme::AutoDiffOp op,
                                PatternRewriter &rewriter) const override {
    SymbolTableCollection symbolTable;
    if (!mlir::enzyme::inlineAutodiffOp(op, rewriter, symbolTable))
      return failure();
 
    return success();
  }
};
 
struct InlineEnzymeForwardDiff
    : public OpRewritePattern<enzyme::ForwardDiffOp> {
  using OpRewritePattern<enzyme::ForwardDiffOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(enzyme::ForwardDiffOp op,
                                PatternRewriter &rewriter) const override {
    SymbolTableCollection symbolTable;
    FunctionOpInterface fn = dyn_cast_or_null<FunctionOpInterface>(
        symbolTable.lookupNearestSymbolFrom(op, op.getFnAttr()));
 
    if (!fn)
      return failure();
 
    Region &targetRegion = fn.getFunctionBody();
 
    if (targetRegion.empty())
      return failure();
 
    // Use a StringAttr rather than a SymbolRefAttr so the function can get
    // symbol-DCE'd
    auto fnAttr = StringAttr::get(op.getContext(), op.getFn());
    auto regionOp = rewriter.replaceOpWithNewOp<enzyme::ForwardDiffRegionOp>(
        op, op.getResultTypes(), op.getInputs(), op.getActivity(),
        op.getRetActivity(), op.getWidth(), op.getStrongZero(), fnAttr);
 
    serializeFunctionAttributes(fn, regionOp);
    rewriter.cloneRegionBefore(targetRegion, regionOp.getBody(),
                               regionOp.getBody().begin());
 
    SmallVector<Operation *> toErase;
    for (Operation &bodyOp : regionOp.getBody().getOps()) {
      if (bodyOp.hasTrait<OpTrait::ReturnLike>()) {
        PatternRewriter::InsertionGuard insertionGuard(rewriter);
        rewriter.setInsertionPoint(&bodyOp);
        enzyme::YieldOp::create(rewriter, bodyOp.getLoc(),
                                bodyOp.getOperands());
        toErase.push_back(&bodyOp);
      }
    }
 
    for (Operation *opToErase : toErase)
      rewriter.eraseOp(opToErase);
    return success();
  }
};
 
// Based on
// https://github.com/llvm/llvm-project/blob/665da0a1649814471739c41a702e0e9447316b20/mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp
 
/// Identifies operations that are beneficial to sink into kernels. These
/// operations may not have side-effects, as otherwise sinking (and hence
/// duplicating them) is not legal.
static bool isLikelyAnIndexComputation(Operation *op) {
  return matchPattern(op, m_Constant()) ||
         isa<memref::DimOp, arith::SelectOp, arith::CmpIOp>(op);
}
 
/// For a given operation `op`, computes whether it is beneficial to sink the
/// operation into the kernel. An operation can be sunk if doing so does not
/// introduce new kernel arguments. Whether a value is already available in the
/// kernel (and hence does not introduce new arguments) is checked by
/// querying `existingDependencies` and `availableValues`.
/// If an operand is not yet available, we recursively check whether it can be
/// made available by siking its defining op.
/// Operations that are indentified for sinking are added to `beneficiaryOps` in
/// the order they should appear in the kernel. Furthermore, `availableValues`
/// is updated with results that will be available after sinking the identified
/// ops.
static bool extractBeneficiaryOps(
    Operation *op, const SetVector<Value> &existingDependencies,
    SetVector<Operation *> &beneficiaryOps,
    llvm::SmallPtrSetImpl<Value> &availableValues,
    llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
  if (beneficiaryOps.count(op))
    return true;
 
  if (!isSinkingBeneficiary(op))
    return false;
 
  for (Value operand : op->getOperands()) {
    // It is already visible in the kernel, keep going.
    if (availableValues.count(operand))
      continue;
    // Else check whether it can be made available via sinking or already is a
    // dependency.
    Operation *definingOp = operand.getDefiningOp();
    if ((!definingOp || !extractBeneficiaryOps(definingOp, existingDependencies,
                                               beneficiaryOps, availableValues,
                                               isSinkingBeneficiary)) &&
        !existingDependencies.count(operand))
      return false;
  }
  // We will sink the operation, mark its results as now available.
  beneficiaryOps.insert(op);
  for (Value result : op->getResults())
    availableValues.insert(result);
  return true;
}
 
template <typename DiffRegionOp>
static LogicalResult sinkOperationsIntoRegionOp(
    DiffRegionOp regionOp,
    function_ref<bool(Operation *)> isSinkingBeneficiary) {
  assert(isSinkingBeneficiary);
  Region &body = regionOp.getBody();
 
  SetVector<Value> sinkCandidates;
  getUsedValuesDefinedAbove(body, sinkCandidates);
 
  SetVector<Operation *> toBeSunk;
  llvm::SmallPtrSet<Value, 4> availableValues;
  for (Value operand : sinkCandidates) {
    Operation *operandOp = operand.getDefiningOp();
    if (!operandOp)
      continue;
    extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues,
                          isSinkingBeneficiary);
  }
 
  // Insert operations so that the defs get cloned before uses.
  IRMapping map;
  OpBuilder builder(body);
  for (Operation *op : toBeSunk) {
    Operation *clonedOp = builder.clone(*op, map);
    // Only replace uses within the launch op.
    for (auto pair : llvm::zip(op->getResults(), clonedOp->getResults()))
      replaceAllUsesInRegionWith(std::get<0>(pair), std::get<1>(pair), body);
  }
  return success();
}
 
template <typename DiffRegionOp>
static FailureOr<func::FuncOp> outlineAutoDiffFunc(
    DiffRegionOp op, StringRef funcName, SmallVectorImpl<Value> &inputs,
    SmallVectorImpl<enzyme::Activity> &argActivities, OpBuilder &builder) {
  Region &autodiffRegion = op.getBody();
  if (failed(sinkOperationsIntoRegionOp(op, isLikelyAnIndexComputation))) {
    return failure();
  }
 
  SmallVector<Type> argTypes(autodiffRegion.getArgumentTypes()), resultTypes;
  SmallVector<Location> argLocs(autodiffRegion.getNumArguments(), op.getLoc());
  // Infer the result types from an enzyme.yield op
  bool found = false;
  autodiffRegion.walk([&](enzyme::YieldOp yieldOp) {
    found = true;
    llvm::append_range(resultTypes, yieldOp.getOperandTypes());
    return WalkResult::interrupt();
  });
  if (!found)
    return op.emitError()
           << "enzyme.yield was not found in enzyme.autodiff_region";
 
  llvm::SetVector<Value> freeValues;
  getUsedValuesDefinedAbove(autodiffRegion, freeValues);
 
  llvm::SmallVector<Value> primalValuesAbove = op.getPrimalInputs();
  llvm::SmallVector<Value> blockArgs(autodiffRegion.getArguments());
  for (Value value : llvm::make_early_inc_range(freeValues)) {
    bool isPrimal = false;
    for (auto [pval, bval] : llvm::zip(primalValuesAbove, blockArgs)) {
      if (value == pval) {
        isPrimal = true;
        for (OpOperand &use : llvm::make_early_inc_range(value.getUses())) {
          if (op->isProperAncestor(use.getOwner()))
            use.assign(bval);
        }
        freeValues.remove(value);
      }
    }
 
    if (!isPrimal) {
      inputs.push_back(value);
      argTypes.push_back(value.getType());
      argLocs.push_back(value.getLoc());
      argActivities.push_back(enzyme::Activity::enzyme_const);
    }
  }
  auto fnType = builder.getFunctionType(argTypes, resultTypes);
 
  // FIXME: making this location the location of the
  // enzyme.autodiff_region op causes translation to LLVM IR to fail due
  // to some issue with the dbg info.
  Location loc = UnknownLoc::get(op.getContext());
  auto outlinedFunc = func::FuncOp::create(builder, loc, funcName, fnType);
  outlinedFunc.setPrivate();
  Region &outlinedBody = outlinedFunc.getBody();
  deserializeFunctionAttributes(op, outlinedFunc, freeValues.size());
 
  // Copy over the function body.
  IRMapping map;
  Block *entryBlock = builder.createBlock(&outlinedBody, outlinedBody.begin(),
                                          argTypes, argLocs);
  unsigned originalArgCount = autodiffRegion.getNumArguments();
  for (const auto &arg : autodiffRegion.getArguments())
    map.map(arg, entryBlock->getArgument(arg.getArgNumber()));
  for (const auto &operand : enumerate(freeValues))
    map.map(operand.value(),
            entryBlock->getArgument(originalArgCount + operand.index()));
  autodiffRegion.cloneInto(&outlinedBody, map);
 
  // Replace the terminators with returns
  for (Block &block : autodiffRegion) {
    Block *clonedBlock = map.lookup(&block);
    auto terminator = dyn_cast<enzyme::YieldOp>(clonedBlock->getTerminator());
    if (!terminator)
      continue;
    OpBuilder replacer(terminator);
    func::ReturnOp::create(replacer, terminator->getLoc(),
                           terminator->getOperands());
    terminator->erase();
  }
 
  // cloneInto results in two blocks, the actual outlined entry block and the
  // cloned autodiff_region entry block. Splice the cloned entry block into
  // the actual entry block, then erase the cloned autodiff_region entry.
  Block *clonedEntry = map.lookup(&autodiffRegion.front());
  entryBlock->getOperations().splice(entryBlock->getOperations().end(),
                                     clonedEntry->getOperations());
  clonedEntry->erase();
  return outlinedFunc;
}
 
LogicalResult outlineEnzymeAutoDiffRegion(enzyme::AutoDiffRegionOp op,
                                          StringRef funcName,
                                          OpBuilder &builder) {
  OpBuilder::InsertionGuard insertionGuard(builder);
  builder.setInsertionPointAfter(op->getParentOfType<SymbolOpInterface>());
 
  SmallVector<enzyme::Activity> argActivities =
      llvm::map_to_vector(op.getActivity().getAsRange<enzyme::ActivityAttr>(),
                          [](auto attr) { return attr.getValue(); });
  size_t numSeeds = llvm::count_if(
      op.getRetActivity().getAsRange<enzyme::ActivityAttr>(), [](auto attr) {
        switch (attr.getValue()) {
        case enzyme::Activity::enzyme_active:
        case enzyme::Activity::enzyme_activenoneed:
          return true;
        default:
          return false;
        }
      });
  size_t numPrimalsAndShadows = op.getInputs().size() - numSeeds;
  SmallVector<Value> primalsAndShadows(
      op.getInputs().begin(), op.getInputs().begin() + numPrimalsAndShadows);
  SmallVector<Value> seedInputs(op.getInputs().begin() + numPrimalsAndShadows,
                                op.getInputs().end());
 
  // Free variables are appended to primalInputs.
  // The final input ordering should be:
  // 1. primals and duplicated argument shadows
  // 2. free variables
  // 3. return variable shadows (a.k.a. seeds)
  FailureOr<func::FuncOp> outlinedFunc = outlineAutoDiffFunc(
      op, funcName, primalsAndShadows, argActivities, builder);
  if (failed(outlinedFunc))
    return failure();
 
  SmallVector<Value> allInputs;
  allInputs.append(primalsAndShadows);
  allInputs.append(seedInputs);
 
  builder.setInsertionPoint(op);
  ArrayAttr argActivityAttr = builder.getArrayAttr(llvm::map_to_vector(
      argActivities, [&op](enzyme::Activity actv) -> Attribute {
        return enzyme::ActivityAttr::get(op.getContext(), actv);
      }));
  auto newOp = enzyme::AutoDiffOp::create(
      builder, op.getLoc(), op.getResultTypes(), outlinedFunc->getName(),
      allInputs, argActivityAttr, op.getRetActivity(), op.getWidth(),
      op.getStrongZero(), op.getAtomicAdd());
  op.replaceAllUsesWith(newOp.getResults());
  op.erase();
  return success();
}
 
LogicalResult outlineEnzymeForwardDiffRegion(enzyme::ForwardDiffRegionOp op,
                                             StringRef funcName,
                                             OpBuilder &builder) {
  OpBuilder::InsertionGuard insertionGuard(builder);
  builder.setInsertionPointAfter(op->getParentOfType<SymbolOpInterface>());
 
  SmallVector<enzyme::Activity> argActivities =
      llvm::map_to_vector(op.getActivity().getAsRange<enzyme::ActivityAttr>(),
                          [](auto attr) { return attr.getValue(); });
  SmallVector<Value> primalsAndShadows(op.getInputs());
 
  // Free variables are appended to primalInputs.
  // The final input ordering should be:
  // 1. primals and duplicated argument shadows
  // 2. free variables
  FailureOr<func::FuncOp> outlinedFunc = outlineAutoDiffFunc(
      op, funcName, primalsAndShadows, argActivities, builder);
  if (failed(outlinedFunc))
    return failure();
 
  SmallVector<Value> allInputs;
  allInputs.append(primalsAndShadows);
 
  builder.setInsertionPoint(op);
  ArrayAttr argActivityAttr = builder.getArrayAttr(llvm::map_to_vector(
      argActivities, [&op](enzyme::Activity actv) -> Attribute {
        return enzyme::ActivityAttr::get(op.getContext(), actv);
      }));
  auto newOp = enzyme::ForwardDiffOp::create(
      builder, op.getLoc(), op.getResultTypes(), outlinedFunc->getName(),
      primalsAndShadows, argActivityAttr, op.getRetActivity(), op.getWidth(),
      op.getStrongZero());
  op.replaceAllUsesWith(newOp.getResults());
  op.erase();
  return success();
}
 
struct InlineEnzymeIntoRegion
    : public enzyme::impl::InlineEnzymeIntoRegionPassBase<
          InlineEnzymeIntoRegion> {
  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.insert<InlineEnzymeAutoDiff, InlineEnzymeForwardDiff>(
        &getContext());
 
    GreedyRewriteConfig config;
    config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal);
    (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
  }
};
 
struct OutlineEnzymeFromRegion
    : public enzyme::impl::OutlineEnzymeFromRegionPassBase<
          OutlineEnzymeFromRegion> {
  void runOnOperation() override {
    SmallVector<enzyme::AutoDiffRegionOp> toOutlineRev;
    getOperation()->walk(
        [&](enzyme::AutoDiffRegionOp op) { toOutlineRev.push_back(op); });
 
    SmallVector<enzyme::ForwardDiffRegionOp> toOutlineFwd;
    getOperation()->walk(
        [&](enzyme::ForwardDiffRegionOp op) { toOutlineFwd.push_back(op); });
 
    OpBuilder builder(getOperation());
    unsigned increment = 0;
    for (auto regionOp : toOutlineRev) {
      auto symbol = regionOp->getParentOfType<SymbolOpInterface>();
      std::string defaultName =
          (Twine(symbol.getName(), "_to_diff") + Twine(increment)).str();
      if (failed(outlineEnzymeAutoDiffRegion(regionOp, defaultName, builder)))
        return signalPassFailure();
 
      ++increment;
    }
 
    for (auto regionOp : toOutlineFwd) {
      auto symbol = regionOp->getParentOfType<SymbolOpInterface>();
      std::string defaultName =
          (Twine(symbol.getName(), "_to_fwddiff") + Twine(increment)).str();
      if (failed(
              outlineEnzymeForwardDiffRegion(regionOp, defaultName, builder)))
        return signalPassFailure();
 
      ++increment;
    }
  }
};
 
} // namespace
 
bool mlir::enzyme::inlineAutodiffOp(enzyme::AutoDiffOp &op,
                                    RewriterBase &rewriter,
                                    SymbolTableCollection &symbolTable) {
  FunctionOpInterface fn = dyn_cast_or_null<FunctionOpInterface>(
      symbolTable.lookupNearestSymbolFrom(op, op.getFnAttr()));
  if (!fn)
    return false;
  Region &targetRegion = fn.getFunctionBody();
  if (targetRegion.empty())
    return false;
 
  // Use a StringAttr rather than a SymbolRefAttr so the function can get
  // symbol-DCE'd
  auto fnAttr = StringAttr::get(op.getContext(), op.getFn());
  auto regionOp = rewriter.replaceOpWithNewOp<enzyme::AutoDiffRegionOp>(
      op, op.getResultTypes(), op.getInputs(), op.getActivity(),
      op.getRetActivity(), op.getWidth(), op.getStrongZero(), op.getAtomicAdd(),
      fnAttr);
  serializeFunctionAttributes(fn, regionOp);
  rewriter.cloneRegionBefore(targetRegion, regionOp.getBody(),
                             regionOp.getBody().begin());
  SmallVector<Operation *> toErase;
  for (Operation &bodyOp : regionOp.getBody().getOps()) {
    if (bodyOp.hasTrait<OpTrait::ReturnLike>()) {
      PatternRewriter::InsertionGuard insertionGuard(rewriter);
      rewriter.setInsertionPoint(&bodyOp);
      enzyme::YieldOp::create(rewriter, bodyOp.getLoc(), bodyOp.getOperands());
      toErase.push_back(&bodyOp);
    }
  }
 
  for (Operation *opToErase : toErase)
    rewriter.eraseOp(opToErase);
  return true;
}