EnzymeToMemRef.cpp

Go to the documentation of this file.

//===- EnzymeToMemRef.cpp - Lower custom Enzyme operations ------------------ //
//
// 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 pass to lower custom ops generated by the Enzyme AD
// procedure to the MemRef dialect.
//===----------------------------------------------------------------------===//
 
#include "Dialect/Dialect.h"
#include "Dialect/Ops.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/Support/raw_ostream.h"
 
namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_ENZYMEOPSTOMEMREFPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
using namespace mlir;
using llvm::errs;
namespace {
struct LoweredCache {
  Value elements, size, capacity;
  Type elementType;
 
  void allocateCache(Location loc, OpBuilder &b) {}
 
  FlatSymbolRefAttr getOrInsertPushFunction(Location loc, ModuleOp moduleOp,
                                            OpBuilder &b) const {
    MLIRContext *context = b.getContext();
    std::string funcName = "__enzyme_push_";
    llvm::raw_string_ostream funcStream{funcName};
    funcStream << elementType;
    if (moduleOp.lookupSymbol<func::FuncOp>(funcName)) {
      return SymbolRefAttr::get(context, funcName);
    }
 
    OpBuilder::InsertionGuard insertionGuard(b);
    b.setInsertionPointToStart(moduleOp.getBody());
 
    auto pushFnType = FunctionType::get(
        context, /*inputs=*/
        {elements.getType(), size.getType(), capacity.getType(), elementType},
        /*outputs=*/{});
    auto pushFn = func::FuncOp::create(b, loc, funcName, pushFnType);
    pushFn.setPrivate();
    Block *entryBlock = pushFn.addEntryBlock();
    b.setInsertionPointToStart(entryBlock);
    BlockArgument elementsField = pushFn.getArgument(0);
    BlockArgument sizeField = pushFn.getArgument(1);
    BlockArgument capacityField = pushFn.getArgument(2);
    BlockArgument value = pushFn.getArgument(3);
 
    Value sizeVal = memref::LoadOp::create(b, loc, sizeField);
    Value capacityVal = memref::LoadOp::create(b, loc, capacityField);
 
    Value predicate = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::eq,
                                            sizeVal, capacityVal);
    scf::IfOp::create(
        b, loc, predicate, [&](OpBuilder &thenBuilder, Location loc) {
          Value two = arith::ConstantIndexOp::create(thenBuilder, loc, 2);
          Value newCapacity =
              arith::MulIOp::create(thenBuilder, loc, capacityVal, two);
          Value oldElements =
              memref::LoadOp::create(thenBuilder, loc, elementsField);
          Value newElements = memref::AllocOp::create(
              thenBuilder, loc, cast<MemRefType>(oldElements.getType()),
              newCapacity);
          memref::CopyOp::create(thenBuilder, loc, oldElements, newElements);
          memref::DeallocOp::create(thenBuilder, loc, oldElements);
          memref::StoreOp::create(thenBuilder, loc, newElements, elementsField);
          memref::StoreOp::create(thenBuilder, loc, newCapacity, capacityField);
          scf::YieldOp::create(thenBuilder, loc);
        });
 
    Value elementsVal = memref::LoadOp::create(b, loc, elementsField);
    memref::StoreOp::create(b, loc, value, elementsVal,
                            /*indices=*/sizeVal);
 
    Value one = arith::ConstantIndexOp::create(b, loc, 1);
    Value newSize = arith::AddIOp::create(b, loc, sizeVal, one);
 
    memref::StoreOp::create(b, loc, newSize, sizeField);
    func::ReturnOp::create(b, loc);
 
    return SymbolRefAttr::get(context, funcName);
  }
 
  FlatSymbolRefAttr getOrInsertPopFunction(Location loc, ModuleOp moduleOp,
                                           OpBuilder &b) {
    MLIRContext *context = b.getContext();
    std::string funcName = "__enzyme_pop_";
    llvm::raw_string_ostream funcStream{funcName};
    funcStream << elementType;
    if (moduleOp.lookupSymbol<func::FuncOp>(funcName)) {
      return SymbolRefAttr::get(context, funcName);
    }
 
    OpBuilder::InsertionGuard insertionGuard(b);
    b.setInsertionPointToStart(moduleOp.getBody());
    auto popFnType = FunctionType::get(
        context,
        /*inputs=*/{elements.getType(), size.getType(), capacity.getType()},
        /*outputs=*/elementType);
    auto popFn = func::FuncOp::create(b, loc, funcName, popFnType);
    popFn.setPrivate();
    Block *entryBlock = popFn.addEntryBlock();
    b.setInsertionPointToStart(entryBlock);
    BlockArgument elementsField = popFn.getArgument(0);
    BlockArgument sizeField = popFn.getArgument(1);
 
    Value elementsVal = memref::LoadOp::create(b, loc, elementsField);
    Value sizeVal = memref::LoadOp::create(b, loc, sizeField);
    Value zero = arith::ConstantIndexOp::create(b, loc, 0);
    Value pred =
        arith::CmpIOp::create(b, loc, arith::CmpIPredicate::sgt, sizeVal, zero);
    cf::AssertOp::create(b, loc, pred, "pop on empty cache");
 
    Value one = arith::ConstantIndexOp::create(b, loc, 1);
    Value newSize = arith::SubIOp::create(b, loc, sizeVal, one);
    memref::StoreOp::create(b, loc, newSize, sizeField);
 
    Value result = memref::LoadOp::create(b, loc, elementsVal, newSize);
    func::ReturnOp::create(b, loc, result);
    return SymbolRefAttr::get(context, funcName);
  }
 
  void emitPush(Location loc, Value value, OpBuilder &b,
                FlatSymbolRefAttr pushFn) const {
    func::CallOp::create(
        b, loc, pushFn, /*results=*/TypeRange{},
        /*operands=*/ValueRange{elements, size, capacity, value});
  }
 
  Value emitPop(Location loc, OpBuilder &b, FlatSymbolRefAttr popFn) const {
    return func::CallOp::create(
               b, loc, popFn, /*results=*/
               cast<ShapedType>(
                   cast<ShapedType>(elements.getType()).getElementType())
                   .getElementType(),
               ValueRange{elements, size, capacity})
        .getResult(0);
  }
 
  FlatSymbolRefAttr getOrInsertGetFunction(Location loc, ModuleOp moduleOp,
                                           OpBuilder &b) {
    MLIRContext *context = b.getContext();
    std::string funcName = "__enzyme_get_";
    llvm::raw_string_ostream funcStream{funcName};
    funcStream << elementType;
    if (moduleOp.lookupSymbol<func::FuncOp>(funcName)) {
      return SymbolRefAttr::get(context, funcName);
    }
 
    OpBuilder::InsertionGuard insertionGuard(b);
    b.setInsertionPointToStart(moduleOp.getBody());
    auto popFnType = FunctionType::get(
        context,
        /*inputs=*/{elements.getType(), size.getType(), capacity.getType()},
        /*outputs=*/elementType);
    auto popFn = func::FuncOp::create(b, loc, funcName, popFnType);
    popFn.setPrivate();
    Block *entryBlock = popFn.addEntryBlock();
    b.setInsertionPointToStart(entryBlock);
    BlockArgument elementsField = popFn.getArgument(0);
    BlockArgument sizeField = popFn.getArgument(1);
 
    Value elementsVal = memref::LoadOp::create(b, loc, elementsField);
    Value sizeVal = memref::LoadOp::create(b, loc, sizeField);
    Value zero = arith::ConstantIndexOp::create(b, loc, 0);
    Value one = arith::ConstantIndexOp::create(b, loc, 1);
    Value pred =
        arith::CmpIOp::create(b, loc, arith::CmpIPredicate::sgt, sizeVal, zero);
    cf::AssertOp::create(b, loc, pred, "get on empty cache");
 
    Value lastIndex = arith::SubIOp::create(b, loc, sizeVal, one);
    Value result = memref::LoadOp::create(b, loc, elementsVal, lastIndex);
    func::ReturnOp::create(b, loc, result);
    return SymbolRefAttr::get(context, funcName);
  }
 
  Value emitGet(Location loc, OpBuilder &b, FlatSymbolRefAttr getFn) const {
    return func::CallOp::create(
               b, loc, getFn, /*results=*/
               cast<ShapedType>(
                   cast<ShapedType>(elements.getType()).getElementType())
                   .getElementType(),
               ValueRange{elements, size, capacity})
        .getResult(0);
  }
  static std::optional<LoweredCache>
  getFromEnzymeCache(Location loc, const TypeConverter *typeConverter,
                     Value enzymeCache, OpBuilder &b) {
    assert(isa<enzyme::CacheType>(enzymeCache.getType()));
    auto cacheType = cast<enzyme::CacheType>(enzymeCache.getType());
    SmallVector<Type> resultTypes;
    if (failed(typeConverter->convertType(cacheType, resultTypes))) {
      return {};
    }
    auto unpackedCache =
        UnrealizedConversionCastOp::create(b, loc, resultTypes, enzymeCache);
    return LoweredCache{.elements = unpackedCache.getResult(0),
                        .size = unpackedCache.getResult(1),
                        .capacity = unpackedCache.getResult(2),
                        .elementType = cacheType.getType()};
  }
};
 
struct InitOpConversion : public OpConversionPattern<enzyme::InitOp> {
  using OpConversionPattern::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(enzyme::InitOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // `enzyme.init` is overloaded to initialize both gradients and caches.
    // Gradients lower to single element MemRefs, while caches lower to
    // variable-sized MemRefs.
 
    if (isa<enzyme::GradientType>(op.getType())) {
      auto memrefType =
          cast<MemRefType>(getTypeConverter()->convertType(op.getType()));
      Value buffer = memref::AllocOp::create(rewriter, op.getLoc(), memrefType);
      rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(op, op.getType(),
                                                              buffer);
      return success();
    }
 
    if (isa<enzyme::CacheType>(op.getType())) {
      SmallVector<Type> resultTypes;
      if (failed(getTypeConverter()->convertType(op.getType(), resultTypes))) {
        op.emitError() << "Failed to convert type " << op.getType();
        return failure();
      }
 
      Value capacity = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 1);
      Value initialSize =
          arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
      auto dataType = cast<MemRefType>(resultTypes[0]);
      auto sizeType = cast<MemRefType>(resultTypes[1]);
      auto capacityType = cast<MemRefType>(resultTypes[2]);
      Value buffer = memref::AllocOp::create(
          rewriter, op.getLoc(), cast<MemRefType>(dataType.getElementType()),
          /*dynamicSize=*/capacity);
      Value bufferField =
          memref::AllocaOp::create(rewriter, op.getLoc(), dataType);
      Value sizeField =
          memref::AllocaOp::create(rewriter, op.getLoc(), sizeType);
      Value capacityField =
          memref::AllocaOp::create(rewriter, op.getLoc(), capacityType);
      memref::StoreOp::create(rewriter, op.getLoc(), buffer, bufferField);
      memref::StoreOp::create(rewriter, op.getLoc(), initialSize, sizeField);
      memref::StoreOp::create(rewriter, op.getLoc(), capacity, capacityField);
      rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
          op, op.getType(), ValueRange{bufferField, sizeField, capacityField});
      return success();
    }
 
    // TODO: Add verification to the init op to verify the valid types,
    // or break out init gradient semantics from init cache semantics
    op.emitError() << "Expected cache or gradient type but got: "
                   << op.getType();
    return failure();
  }
};
 
struct PushOpConversion : public OpConversionPattern<enzyme::PushOp> {
  using OpConversionPattern::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(enzyme::PushOp op, OneToNOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto loweredCache = LoweredCache::getFromEnzymeCache(
        loc, getTypeConverter(), op.getCache(), rewriter);
    if (!loweredCache.has_value()) {
      return failure();
    }
 
    FlatSymbolRefAttr pushFn = loweredCache.value().getOrInsertPushFunction(
        loc, op->getParentOfType<ModuleOp>(), rewriter);
    loweredCache.value().emitPush(loc, op.getValue(), rewriter, pushFn);
    rewriter.eraseOp(op);
    return success();
  }
};
 
struct PopOpConversion : public OpConversionPattern<enzyme::PopOp> {
  using OpConversionPattern::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(enzyme::PopOp op, OneToNOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto loweredCache = LoweredCache::getFromEnzymeCache(
        loc, getTypeConverter(), op.getCache(), rewriter);
    if (!loweredCache.has_value()) {
      return failure();
    }
 
    FlatSymbolRefAttr popFn = loweredCache.value().getOrInsertPopFunction(
        loc, op->getParentOfType<ModuleOp>(), rewriter);
    rewriter.replaceOp(op, loweredCache.value().emitPop(loc, rewriter, popFn));
    return success();
  }
};
 
struct SetOpConversion : public OpConversionPattern<enzyme::SetOp> {
  using OpConversionPattern::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(enzyme::SetOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (auto type = dyn_cast<enzyme::CacheType>(op.getGradient().getType())) {
      op.emitError() << "set for CacheType not implemented";
      return failure();
    } else if (auto type =
                   dyn_cast<enzyme::GradientType>(op.getGradient().getType())) {
      auto memrefType = cast<MemRefType>(getTypeConverter()->convertType(type));
      auto castedGradient = UnrealizedConversionCastOp::create(
          rewriter, op.getLoc(), memrefType, op.getGradient());
      rewriter.replaceOpWithNewOp<memref::StoreOp>(op, op.getValue(),
                                                   castedGradient.getResult(0));
    }
    return success();
  }
};
 
struct GetOpConversion : public OpConversionPattern<enzyme::GetOp> {
  using OpConversionPattern<enzyme::GetOp>::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(enzyme::GetOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
 
    if (auto type = dyn_cast<enzyme::CacheType>(op.getGradient().getType())) {
      Location loc = op.getLoc();
      auto loweredCache = LoweredCache::getFromEnzymeCache(
          loc, getTypeConverter(), op.getGradient(), rewriter);
      if (!loweredCache.has_value()) {
        return failure();
      }
 
      FlatSymbolRefAttr getFn = loweredCache.value().getOrInsertGetFunction(
          loc, op->getParentOfType<ModuleOp>(), rewriter);
      rewriter.replaceOp(op,
                         loweredCache.value().emitGet(loc, rewriter, getFn));
    } else if (auto type =
                   dyn_cast<enzyme::GradientType>(op.getGradient().getType())) {
      auto memrefType = cast<MemRefType>(getTypeConverter()->convertType(type));
      auto castedGradient = UnrealizedConversionCastOp::create(
          rewriter, op.getLoc(), memrefType, op.getGradient());
      rewriter.replaceOpWithNewOp<memref::LoadOp>(op,
                                                  castedGradient.getResult(0));
    }
    return success();
  }
};
 
struct EnzymeToMemRefPass
    : public enzyme::impl::EnzymeOpsToMemRefPassBase<EnzymeToMemRefPass> {
  void runOnOperation() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    TypeConverter typeConverter;
    typeConverter.addConversion([](Type type) -> std::optional<Type> {
      if (type.isIntOrIndexOrFloat() || isa<MemRefType>(type))
        return type;
      return {};
    });
    typeConverter.addConversion(
        [](enzyme::GradientType type) -> std::optional<Type> {
          return MemRefType::get({}, type.getBasetype());
        });
    typeConverter.addConversion(
        [](enzyme::CacheType type, SmallVectorImpl<Type> &resultTypes) {
          // Data
          resultTypes.push_back(MemRefType::get(
              {}, MemRefType::get({ShapedType::kDynamic}, type.getType())));
          auto indexMemRefType =
              MemRefType::get({}, IndexType::get(type.getContext()));
          // Size
          resultTypes.push_back(indexMemRefType);
          // Capacity
          resultTypes.push_back(indexMemRefType);
          return success();
        });
 
    patterns.add<InitOpConversion>(typeConverter, context);
    patterns.add<PushOpConversion>(typeConverter, context);
    patterns.add<PopOpConversion>(typeConverter, context);
    patterns.add<SetOpConversion>(typeConverter, context);
    patterns.add<GetOpConversion>(typeConverter, context);
 
    ConversionTarget target(*context);
    target.addLegalDialect<memref::MemRefDialect>();
    target.addLegalDialect<arith::ArithDialect>();
    target.addLegalDialect<scf::SCFDialect>();
    target.addLegalDialect<cf::ControlFlowDialect>();
    target.addLegalDialect<func::FuncDialect>();
    target.addLegalOp<UnrealizedConversionCastOp>();
    target.addIllegalDialect<enzyme::EnzymeDialect>();
 
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  };
};
} // end anonymous namespace