EnzymeBatchPass.h

Go to the documentation of this file.

#ifndef ENZYME_BATCH_PASS_H
#define ENZYME_BATCH_PASS_H
 
#include "Dialect/Ops.h"
#include "Interfaces/GradientUtilsReverse.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
 
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
 
namespace mlir {
namespace enzyme {
namespace batchutils {
 
struct BatchCacheKey {
  FunctionOpInterface function;
  SmallVector<int64_t> batchSizes;
 
  bool operator<(const BatchCacheKey &other) const {
    Operation *thisSymbolTable = SymbolTable::getNearestSymbolTable(function),
              *otherSymbolTable =
                  SymbolTable::getNearestSymbolTable(other.function);
    if (thisSymbolTable != otherSymbolTable)
      return std::less<Operation *>()(thisSymbolTable, otherSymbolTable);
 
    auto thisName = const_cast<FunctionOpInterface &>(function).getName(),
         otherName =
             const_cast<FunctionOpInterface &>(other.function).getName();
    if (thisName != otherName)
      return thisName < otherName;
 
    return std::lexicographical_compare(batchSizes.begin(), batchSizes.end(),
                                        other.batchSizes.begin(),
                                        other.batchSizes.end());
  }
};
 
mlir::TensorType applyBatchSizes(mlir::Type Ty,
                                 llvm::ArrayRef<int64_t> batchSizes);
 
FunctionOpInterface batchCloneFunction(
    OpBuilder &builder, FunctionOpInterface F, Twine name,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache);
 
void batchCloneBlock(
    OpBuilder &builder, Block *blk, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache,
    bool withoutTerminator);
 
void batchCloneRegion(
    OpBuilder &builder, Region *src, Region *dest, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache);
 
LogicalResult handleCallOp(
    func::CallOp callOp, OpBuilder &builder, IRMapping &mapper,
    llvm::ArrayRef<int64_t> batchSizes,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache);
 
template <typename T>
FunctionOpInterface batchOperationWithoutInsertingCallOp(
    OpBuilder &builder, T CI, FunctionOpInterface fn,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache);
 
template <typename T>
LogicalResult batchOperation(
    SymbolTableCollection &symbolTable, OpBuilder &builder, T CI,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
 
  auto *symbolOp = symbolTable.lookupNearestSymbolFrom(CI, CI.getFnAttr());
  return batchOperation(builder, CI, cast<FunctionOpInterface>(symbolOp),
                        batchedFunctionCache);
}
 
template <typename T>
LogicalResult batchOperation(
    SymbolTableCollection &symbolTable, PatternRewriter &rewriter, T CI,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  auto *symbolOp = symbolTable.lookupNearestSymbolFrom(CI, CI.getFnAttr());
  return batchOperation(rewriter, CI, cast<FunctionOpInterface>(symbolOp),
                        batchedFunctionCache);
}
 
template <typename T>
LogicalResult batchOperation(
    OpBuilder &builder, T CI, FunctionOpInterface fn,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  auto newFunc = batchOperationWithoutInsertingCallOp(builder, CI, fn,
                                                      batchedFunctionCache);
 
  if (!newFunc)
    return failure();
 
  {
    IRRewriter::InsertionGuard insertGuard(builder);
    builder.setInsertionPoint(CI);
    auto dCI = func::CallOp::create(builder, CI.getLoc(), newFunc.getName(),
                                    newFunc.getResultTypes(), CI.getInputs());
    CI.replaceAllUsesWith(dCI);
    CI->erase();
  }
  return success();
}
 
template <typename T>
LogicalResult batchOperation(
    PatternRewriter &rewriter, T CI, FunctionOpInterface fn,
    std::map<BatchCacheKey, FunctionOpInterface> &batchedFunctionCache) {
  auto newFunc = batchOperationWithoutInsertingCallOp(rewriter, CI, fn,
                                                      batchedFunctionCache);
 
  if (!newFunc)
    return failure();
 
  {
    IRRewriter::InsertionGuard insertGuard(rewriter);
    rewriter.setInsertionPoint(CI);
    rewriter.replaceOpWithNewOp<func::CallOp>(
        CI, newFunc.getName(), newFunc.getResultTypes(), CI.getInputs());
  }
  return success();
}
 
// instead of inserting a call op, we will inline each operation directly
// into the caller
inline void batchOperationInline(PatternRewriter &rewriter,
                                 enzyme::BatchOp batchOp,
                                 FunctionOpInterface func) {
  auto &origRegion = func.getFunctionBody();
  auto &origBlock = origRegion.front();
 
  IRMapping mapper;
  for (int i = 0; i < batchOp->getNumOperands(); i++) {
    mapper.map(origBlock.getArguments()[i], batchOp->getOperand(i));
  }
 
  rewriter.setInsertionPoint(batchOp);
  std::map<BatchCacheKey, FunctionOpInterface> batchedFunctionCache;
  batchCloneBlock(rewriter, &origBlock, mapper, batchOp.getBatchShape(),
                  batchedFunctionCache, true);
 
  auto origTerm = origBlock.getTerminator();
  for (auto [i, operand] : llvm::enumerate(origTerm->getOperands())) {
    auto mappedOperand = mapper.lookup(operand);
    rewriter.replaceAllUsesWith(batchOp->getResult(i), mappedOperand);
  }
  rewriter.eraseOp(batchOp);
  rewriter.eraseOp(func);
}
 
} // namespace batchutils
} // namespace enzyme
} // namespace mlir
 
#endif // ENZYME_BATCH_PASS_H