AddToOpToIndexAndLoad.cpp

Go to the documentation of this file.

//===- AddToOpToIndexAndLoad.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/Ops.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Transforms/DialectConversion.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_ADDTOOPTOINDEXANDLOADPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
using namespace mlir;
using namespace enzyme;
using llvm::errs;
namespace {
 
SmallVector<Value> applyAffineMap(AffineMap aMap, SmallVector<Value> indices,
                                  OpBuilder &builder, Location loc) {
  SmallVector<Value> appliedAffineMap;
  for (unsigned int i = 0; i < aMap.getNumResults(); i++) {
    AffineMap subMap = aMap.getSubMap({i});
    auto mapApplied = affine::AffineApplyOp::create(builder, loc, subMap,
                                                    ValueRange(indices));
    appliedAffineMap.push_back(mapApplied);
  }
  return appliedAffineMap;
}
 
struct AddToOpToIndexAndLoadPass
    : public enzyme::impl::AddToOpToIndexAndLoadPassBase<
          AddToOpToIndexAndLoadPass> {
  void runOnOperation() override {
    getOperation()->walk([&](Operation *op) {
      auto loc = op->getLoc();
      auto enzymeAdjoint = dyn_cast<enzyme::GenericAdjointOp>(op);
      if (!enzymeAdjoint)
        return;
 
      OpBuilder cacheBuilder(enzymeAdjoint);
      auto adjoint = Utils::adjointToGeneric(enzymeAdjoint, cacheBuilder, loc);
 
      // check if adjoint contains a enzyme.addToOp
      Operation *addToOp = nullptr;
      adjoint.walk([&](Operation *op) {
        if (isa<enzyme::AddToOp>(op)) {
          addToOp = op;
        }
      });
      if (!addToOp)
        return;
 
      Operation *terminator = adjoint.getBodyRegion().front().getTerminator();
      SmallVector<Value> indices;
      SmallVector<Value> retargs;
      auto outs = adjoint.getOutputs();
      auto num_ins = adjoint.getInputs().size();
      for (auto val : addToOp->getOperands()) {
        retargs.push_back(val);
      }
      auto map = adjoint.getIndexingMapsArray();
      cacheBuilder.setInsertionPoint(terminator);
 
      // Is it a fine assumption that all indexing maps are the same?
      for (size_t i = 0; i < map[0].getNumDims(); i++) {
        indices.push_back(linalg::IndexOp::create(cacheBuilder, loc, i));
      }
 
      SmallVector<Value> rets;
      for (size_t i = 0; i < retargs.size(); i++) {
        // auto load = AffineLoadOp::create(cacheBuilder, loc, inputs[i],
        // map[i], indices); auto store = AffineStoreOp::create(cacheBuilder,
        // loc, load, inputs[i], map[i], indices);
        SmallVector<Value> mapAppliedIndices =
            applyAffineMap(map[num_ins + i], indices, cacheBuilder, loc);
        auto load = memref::LoadOp::create(cacheBuilder, loc, outs[i],
                                           mapAppliedIndices);
        auto added = cast<enzyme::AutoDiffTypeInterface>(load.getType())
                         .createAddOp(cacheBuilder, loc, load, retargs[i]);
        memref::StoreOp::create(cacheBuilder, loc, added, outs[i],
                                mapAppliedIndices);
      }
 
      for (size_t i = 0; i < retargs.size(); i++) {
        SmallVector<Value> mapAppliedIndices =
            applyAffineMap(map[num_ins + i], indices, cacheBuilder, loc);
        auto load = memref::LoadOp::create(cacheBuilder, loc, outs[i],
                                           mapAppliedIndices);
        retargs[i] = load;
      }
 
      linalg::YieldOp::create(cacheBuilder, loc, ValueRange{retargs});
      addToOp->erase();
    });
  };
};
} // end anonymous namespace