AddToOpToSplit.cpp

Go to the documentation of this file.

//===- AddToOpToSplit.cpp - Lower Shadowed Gradient ops
//------------------ //
//
// 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 "mlir/Rewrite/PatternApplicator.h"
#include "llvm/Support/raw_ostream.h"
 
#include "Interfaces/AutoDiffTypeInterface.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
 
#include "Utils.h"
 
namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_ADDTOOPTOSPLITPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
using namespace mlir;
using namespace enzyme;
using llvm::errs;
namespace {
 
Operation *getAddToOp(linalg::GenericOp &adjoint) {
  Operation *addToOp = nullptr;
  adjoint.walk([&](Operation *op) {
    if (isa<enzyme::AddToOp>(op)) {
      addToOp = op;
    }
  });
  return addToOp;
}
 
bool isMemrefCacheType(Type type) {
  if (auto memrefType = dyn_cast<MemRefType>(type)) {
    return isa<CacheType>(memrefType.getElementType());
  }
  return false;
}
 
void processGenericDuplication(Operation *op, OpBuilder &builder, Location loc,
                               int i) {
  auto clonedAdjoint = builder.clone(*op);
  linalg::GenericOp clonedAdjointGenericOp =
      cast<linalg::GenericOp>(clonedAdjoint);
 
  // Delete all but the ith output
  unsigned numInputs = clonedAdjointGenericOp.getInputsMutable().size();
  SmallVector<mlir::Attribute> indexingMaps(
      clonedAdjointGenericOp.getIndexingMaps().getValue());
 
  if (clonedAdjointGenericOp.getOutputsMutable().size() - i - 1 > 0) {
    unsigned idx = i + 1;
    unsigned len = clonedAdjointGenericOp.getOutputsMutable().size() - i - 1;
    clonedAdjointGenericOp.getOutputsMutable().erase(idx, len);
    clonedAdjointGenericOp.getRegion().front().eraseArguments(numInputs + idx,
                                                              len);
    indexingMaps.erase(indexingMaps.begin() + numInputs + idx,
                       indexingMaps.begin() + numInputs + idx + len);
  }
 
  if (i > 0) {
    clonedAdjointGenericOp.getOutputsMutable().erase(0, i);
    clonedAdjointGenericOp.getRegion().front().eraseArguments(numInputs, i);
    indexingMaps.erase(indexingMaps.begin() + numInputs,
                       indexingMaps.begin() + numInputs + i);
  }
  clonedAdjointGenericOp.setIndexingMapsAttr(
      builder.getArrayAttr(indexingMaps));
 
  auto clonedAddToOp = getAddToOp(clonedAdjointGenericOp);
  auto scope = OpBuilder::InsertionGuard(builder);
 
  builder.setInsertionPointAfter(clonedAddToOp);
  auto terminator = linalg::YieldOp::create(builder, loc);
 
  auto operand = clonedAddToOp->getOperand(i);
  auto outputOperand =
      clonedAdjointGenericOp.getRegion().front().getArgument(numInputs + i);
  auto operandType = cast<AutoDiffTypeInterface>(operand.getType());
 
  builder.setInsertionPoint(terminator);
  auto returnValue =
      operandType.createAddOp(builder, loc, outputOperand, operand);
 
  terminator->setOperands({returnValue});
  clonedAddToOp->erase();
}
 
struct AddToOpToSplitPass
    : public enzyme::impl::AddToOpToSplitPassBase<AddToOpToSplitPass> {
  void runOnOperation() override {
    getOperation()->walk([&](Operation *op) {
      auto enzymeAdjoint = dyn_cast<enzyme::GenericAdjointOp>(op);
      auto loc = op->getLoc();
      if (!enzymeAdjoint)
        return;
 
      OpBuilder builder(enzymeAdjoint);
      auto adjoint = Utils::adjointToGeneric(enzymeAdjoint, builder, loc);
 
      Operation *addToOp = getAddToOp(adjoint);
      if (!addToOp)
        return;
 
      // TODO duplicate memref<CacheType> inputs
      // For now just error out
      for (auto input : adjoint.getInputs()) {
        if (isMemrefCacheType(input.getType())) {
          llvm::report_fatal_error(
              "Cannot split AddToOp with memref<CacheType> inputs");
          return;
        }
      }
 
      builder.setInsertionPoint(adjoint);
 
      for (size_t i = 0; i < addToOp->getNumOperands(); i++) {
        processGenericDuplication(adjoint.getOperation(), builder, loc, i);
      }
      adjoint->erase();
    });
  };
};
} // end anonymous namespace