EnzymeBatchPass.h File Reference

#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"

Include dependency graph for EnzymeBatchPass.h:

This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Classes
struct	mlir::enzyme::batchutils::BatchCacheKey

Namespaces
namespace	mlir
namespace	mlir::enzyme
namespace	mlir::enzyme::batchutils

Functions
mlir::TensorType	mlir::enzyme::batchutils::applyBatchSizes (mlir::Type Ty, llvm::ArrayRef< int64_t > batchSizes)
FunctionOpInterface	mlir::enzyme::batchutils::batchCloneFunction (OpBuilder &builder, FunctionOpInterface F, Twine name, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
void	mlir::enzyme::batchutils::batchCloneBlock (OpBuilder &builder, Block *blk, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache, bool withoutTerminator)
void	mlir::enzyme::batchutils::batchCloneRegion (OpBuilder &builder, Region *src, Region *dest, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
LogicalResult	mlir::enzyme::batchutils::handleCallOp (func::CallOp callOp, OpBuilder &builder, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
FunctionOpInterface	mlir::enzyme::batchutils::batchOperationWithoutInsertingCallOp (OpBuilder &builder, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	mlir::enzyme::batchutils::batchOperation (SymbolTableCollection &symbolTable, OpBuilder &builder, T CI, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	mlir::enzyme::batchutils::batchOperation (SymbolTableCollection &symbolTable, PatternRewriter &rewriter, T CI, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	mlir::enzyme::batchutils::batchOperation (OpBuilder &builder, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	mlir::enzyme::batchutils::batchOperation (PatternRewriter &rewriter, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
void	mlir::enzyme::batchutils::batchOperationInline (PatternRewriter &rewriter, enzyme::BatchOp batchOp, FunctionOpInterface func)