EnzymeBatchDiffPass.cpp

Go to the documentation of this file.

//===- EnzymeBatchDiffPass.cpp - Merge autodiff calls into their batched
// versions
//------------ //
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "Passes/EnzymeBatchDiffPass.h"
#include "Dialect/Ops.h"
#include "Interfaces/GradientUtilsReverse.h"
#include "Interfaces/Utils.h"
#include "PassDetails.h"
#include "Passes/Passes.h"
#include "Passes/Utils.h"
 
#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "llvm/ADT/DenseMap.h"
 
#define DEBUG_TYPE "enzyme-diff-batch"
#define ENZYME_DBGS llvm::dbgs() << "[" << DEBUG_TYPE << "]"
 
using namespace mlir;
using namespace mlir::enzyme;
using namespace enzyme;
 
namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_BATCHDIFFPASS
#include "Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir
 
namespace {
 
struct BatchDiffPass : public enzyme::impl::BatchDiffPassBase<BatchDiffPass> {
  void runOnOperation() override;
 
  void mergeFwddiffCalls(SymbolTableCollection &symbolTable,
                         FunctionOpInterface op) {
    // TODO: Use a modified version of inter-procedural DataFlowAliasAnalysis
    // for mapping primal effects
    llvm::DenseMap<FunctionOpInterface,
                   SmallVector<MemoryEffects::EffectInstance>>
        innerEffectCache;
 
    OpBuilder builder(op);
 
    op->walk([&](Block *blk) {
      std::map<enzyme::batchutils::BatchDiffCacheKey,
               SmallVector<enzyme::ForwardDiffOp>>
          diffMergeSet;
 
      for (auto fwdOp : blk->getOps<enzyme::ForwardDiffOp>()) {
        auto fnOp = dyn_cast_or_null<FunctionOpInterface>(
            symbolTable.lookupNearestSymbolFrom(fwdOp, fwdOp.getFnAttr()));
        if (!fnOp)
          continue;
 
        batchutils::BatchDiffCacheKey key =
            batchutils::createDiffCacheKey(fwdOp, fnOp);
 
        diffMergeSet[key].push_back(fwdOp);
      }
 
      for (auto &pair : diffMergeSet) {
        auto key = pair.first;
        auto allDiffs = pair.second;
        if (allDiffs.size() < 2)
          continue;
 
        // Collect inner effects of function
        if (!innerEffectCache.contains(key.function)) {
          innerEffectCache[key.function] =
              oputils::collectFnEffects(key.function);
        }
 
        SmallVector<MemoryEffects::EffectInstance> &calleeEffects =
            innerEffectCache[key.function];
 
        // TODO: skip if known readnone from existing analyses
        bool skipMergeEntry = false;
 
        // Map callee(primal function) memory effects to the calling
        // function's(autodiff op's) memory effects. This allows us to also
        // reason about memory effects on the derivative argument potentially
        // being passed in(if the primal argument has activity enzyme_dup)
 
        // ForwardDiff read(primal) ?-> read(deriv);
        // write(primal) ?-> write(deriv)
        //
        // ReverseDiff read(primal) ?-> read or write(deriv);
        // write(primal) ?-> read or write(deriv);
        //
        // Unknown effects are retained in the caller, and will always alias to
        // true with any other effect
        llvm::DenseMap<ForwardDiffOp,
                       SmallVector<MemoryEffects::EffectInstance>>
            callerEffectMap;
 
        for (auto &eff : calleeEffects) {
          if (!isa<MemoryEffects::Read>(eff.getEffect()) &&
              !isa<MemoryEffects::Write>(eff.getEffect())) {
            // encountered allocate/load, skip merging
            skipMergeEntry = true;
            break;
          }
 
          Value effVal = eff.getValue();
          if (!effVal) {
            // unknown effect which isnt a value, skip merging
            skipMergeEntry = true;
            break;
          }
 
          // Find primal argument corresponding to effect value
          size_t primalArgPos = 0;
          bool foundPrimal = false;
          if (auto effBA = dyn_cast<BlockArgument>(effVal)) {
            if (llvm::is_contained(key.function.getArguments(), effBA)) {
              foundPrimal = true;
              primalArgPos = effBA.getArgNumber();
            }
          }
 
          if (!foundPrimal) {
            // TODO: Handle this either as a global value, or a value which
            // is inside of the MLIR function(for inter-proc alias analysis) -
            // skip for now
            skipMergeEntry = true;
            break;
          }
 
          // Add primal effects to caller effect map for all ops
          Value primalVal = key.inputs[primalArgPos];
          for (auto dop : allDiffs) {
            callerEffectMap[dop].push_back(oputils::getEffectOfVal(
                primalVal, eff.getEffect(), eff.getResource()));
          }
 
          // Add derivative effects(only if primal arg is dup)
          // read(primal) -> read(derivative)
          // write(primal) -> write(derivative)
 
          // find position of dup arg for primal
          bool primalIsDup =
              (key.inActivity[primalArgPos] == Activity::enzyme_dup) ||
              (key.inActivity[primalArgPos] == Activity::enzyme_dupnoneed);
 
          if (primalIsDup) {
            size_t gradArgPos = 0;
            for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
              ++gradArgPos;
 
              if (idx == primalArgPos)
                break;
 
              if (act == Activity::enzyme_dup ||
                  act == Activity::enzyme_dupnoneed) {
                ++gradArgPos;
              }
            }
 
            for (auto dop : allDiffs) {
              Value dVal = dop.getInputs()[gradArgPos];
              callerEffectMap[dop].push_back(oputils::getEffectOfVal(
                  dVal, eff.getEffect(), eff.getResource()));
            }
          }
        }
 
        if (skipMergeEntry)
          continue;
 
        SmallVector<ForwardDiffOp> prunedSources =
            batchutils::pruneGradDefs(key, allDiffs);
 
        SmallVector<ForwardDiffOp> legalMerge = batchutils::pruneMemoryEffects(
            symbolTable, key, prunedSources, callerEffectMap, innerEffectCache);
 
        {
          SmallVector<enzyme::ForwardDiffOp> &allOps = legalMerge;
          int64_t width = allOps.size();
          if (width < 2)
            continue;
 
          auto firstDiffOp = allOps.front();
          IRRewriter::InsertionGuard insertGuard(builder);
          builder.setInsertionPoint(firstDiffOp);
          auto loc = firstDiffOp->getLoc();
          auto context = builder.getContext();
 
          SmallVector<mlir::Value> in_args;
          SmallVector<ActivityAttr, 2> inActivityAttrs;
          SmallVector<ActivityAttr, 2> retActivityAttrs;
          SmallVector<mlir::Type, 2> out_ty;
          auto in_idx = 0;
          for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
            ActivityAttr iattr = ActivityAttr::get(context, act);
            inActivityAttrs.push_back(iattr);
            in_args.push_back(key.inputs[idx]);
            in_idx++;
            // batched input derivative
            SmallVector<mlir::Value> derivList;
            if (act == Activity::enzyme_dup ||
                act == Activity::enzyme_dupnoneed) {
              for (auto uop : allOps) {
                derivList.push_back(uop.getInputs()[in_idx]);
              }
 
              mlir::Value batchedDeriv =
                  getConcatValue(builder, loc, derivList);
              in_args.push_back(batchedDeriv);
              in_idx++;
            }
          }
 
          auto out_idx = 0;
          for (auto [idx, ract] : llvm::enumerate(key.retActivity)) {
            ActivityAttr iattr = ActivityAttr::get(context, ract);
            retActivityAttrs.push_back(iattr);
            switch (ract) {
            case Activity::enzyme_active: {
              mlir::Value res = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(res.getType());
              ++out_idx;
              break;
            }
            case Activity::enzyme_const: {
              mlir::Value res = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(res.getType());
              ++out_idx;
              break;
            }
            case Activity::enzyme_dupnoneed: {
              // batched output derivative
              mlir::Value dres = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(getConcatType(dres, width));
              ++out_idx;
              break;
            }
            case Activity::enzyme_dup: {
              mlir::Value res = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(res.getType());
              ++out_idx;
              // batched output derivative
              mlir::Value dres = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(getConcatType(dres, width));
              ++out_idx;
              break;
            }
            case Activity::enzyme_constnoneed: {
              break;
            }
            case Activity::enzyme_activenoneed: {
              mlir::Value res = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(res.getType());
              ++out_idx;
              break;
            }
            default:
              llvm_unreachable(
                  "unknown activity value encountered for ret_activity");
            }
          }
 
          ArrayAttr newInActivity = ArrayAttr::get(
              context, llvm::ArrayRef<Attribute>(inActivityAttrs.begin(),
                                                 inActivityAttrs.end()));
          ArrayAttr newRetActivity = ArrayAttr::get(
              context, llvm::ArrayRef<Attribute>(retActivityAttrs.begin(),
                                                 retActivityAttrs.end()));
          IntegerAttr newWidthAttr =
              IntegerAttr::get(firstDiffOp.getWidthAttr().getType(), width);
          auto newDiffOp = ForwardDiffOp::create(
              builder, loc, out_ty, firstDiffOp.getFnAttr(), in_args,
              newInActivity, newRetActivity, newWidthAttr,
              firstDiffOp.getStrongZeroAttr());
 
          // Rename uses
          // out -> primal
          // dout -> derivative
          out_idx = 0;
          for (auto [idx, ract] : llvm::enumerate(key.retActivity)) {
            switch (ract) {
            case Activity::enzyme_constnoneed:
              break;
 
            case Activity::enzyme_active:
            case Activity::enzyme_activenoneed:
            case Activity::enzyme_const: {
              auto new_out = newDiffOp.getOutputs()[out_idx];
              for (auto dop : allOps) {
                dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
              }
              out_idx++;
              break;
            }
 
            case Activity::enzyme_dupnoneed: {
              // batched derivative
              auto batch_dout = newDiffOp.getOutputs()[out_idx];
              for (auto [dop_idx, dop] : llvm::enumerate(allOps)) {
                auto old_dout = dop.getOutputs()[out_idx];
                auto doutTy = old_dout.getType();
                auto new_dout =
                    getExtractValue(builder, loc, doutTy, batch_dout, dop_idx);
 
                old_dout.replaceAllUsesWith(new_dout);
              }
              ++out_idx;
              break;
            }
 
            case Activity::enzyme_dup: {
              mlir::Value new_out = newDiffOp.getOutputs()[out_idx];
              for (ForwardDiffOp dop : allOps) {
                dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
              }
              out_idx++;
 
              // batched derivative
              auto batch_dout = newDiffOp.getOutputs()[out_idx];
              for (auto [dop_idx, dop] : llvm::enumerate(allOps)) {
                auto old_dout = dop.getOutputs()[out_idx];
                auto doutTy = old_dout.getType();
                auto new_dout =
                    getExtractValue(builder, loc, doutTy, batch_dout, dop_idx);
 
                old_dout.replaceAllUsesWith(new_dout);
              }
              ++out_idx;
              break;
            }
            default:
              llvm_unreachable(
                  "unknown activity value encountered for ret_activity");
            }
          }
 
          // erase all old ops
          for (auto dop : allOps) {
            dop->erase();
          }
        }
      }
    }); // block walker
  }
 
  void mergeRevdiffCalls(SymbolTableCollection &symbolTable,
                         FunctionOpInterface op) {
 
    // TODO: Use a modified version of inter-procedural DataFlowAliasAnalysis
    // for mapping primal effects
 
    // list of values read/written to inside fn
    llvm::DenseMap<FunctionOpInterface,
                   SmallVector<MemoryEffects::EffectInstance>>
        innerEffectCache;
 
    OpBuilder builder(op);
 
    op->walk([&](Block *blk) {
      // map tracking batchable AD calls
      std::map<enzyme::batchutils::BatchDiffCacheKey,
               SmallVector<enzyme::AutoDiffOp>>
          diffMergeSet;
 
      for (auto revOp : blk->getOps<enzyme::AutoDiffOp>()) {
        auto fnOp = dyn_cast_or_null<FunctionOpInterface>(
            symbolTable.lookupNearestSymbolFrom(revOp, revOp.getFnAttr()));
        if (!fnOp)
          continue;
 
        batchutils::BatchDiffCacheKey key =
            batchutils::createDiffCacheKey(revOp, fnOp);
 
        diffMergeSet[key].push_back(revOp);
      }
 
      for (auto &pair : diffMergeSet) {
        auto key = pair.first;
        auto allDiffs = pair.second;
        if (allDiffs.size() < 2)
          continue;
 
        // Collect inner effects of function
        if (!innerEffectCache.contains(key.function)) {
          innerEffectCache[key.function] =
              oputils::collectFnEffects(key.function);
        }
 
        SmallVector<MemoryEffects::EffectInstance> &calleeEffects =
            innerEffectCache[key.function];
 
        // TODO: skip if known readonly from existing analyses
        bool skipMergeEntry = false;
 
        llvm::DenseMap<AutoDiffOp, SmallVector<MemoryEffects::EffectInstance>>
            callerEffectMap;
 
        for (auto &eff : calleeEffects) {
          if (!isa<MemoryEffects::Read>(eff.getEffect()) &&
              !isa<MemoryEffects::Write>(eff.getEffect())) {
            // encountered allocate/load, skip merging
            skipMergeEntry = true;
            break;
          }
 
          Value effVal = eff.getValue();
          if (!effVal) {
            // unknown effect which isnt a value, skip merging
            skipMergeEntry = true;
            break;
          }
 
          // Find primal argument corresponding to effect value
          size_t primalArgPos = 0;
          bool foundPrimal = false;
          if (auto effBA = dyn_cast<BlockArgument>(effVal)) {
            if (llvm::is_contained(key.function.getArguments(), effBA)) {
              foundPrimal = true;
              primalArgPos = effBA.getArgNumber();
            }
          }
 
          if (!foundPrimal) {
            // TODO: Handle this either as a global value, or a value which
            // is inside of the MLIR function(for inter-proc alias analysis) -
            // skip for now
            skipMergeEntry = true;
            break;
          }
 
          // Add primal effects to caller effect map for all ops
          Value primalVal = key.inputs[primalArgPos];
          for (auto dop : allDiffs) {
            callerEffectMap[dop].push_back(oputils::getEffectOfVal(
                primalVal, eff.getEffect(), eff.getResource()));
          }
 
          // Add derivative effects(only if primal arg is dup)
          // read(primal) -> read(derivative) + write(derivative)
          // write(primal) -> write(derivative) + read(derivative)
 
          // find position of dup arg for primal
          bool primalIsDup =
              (key.inActivity[primalArgPos] == Activity::enzyme_dup) ||
              (key.inActivity[primalArgPos] == Activity::enzyme_dupnoneed);
 
          if (primalIsDup) {
            size_t gradArgPos = 0;
            for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
              ++gradArgPos;
 
              if (idx == primalArgPos)
                break;
 
              if (act == Activity::enzyme_dup ||
                  act == Activity::enzyme_dupnoneed) {
                ++gradArgPos;
              }
            }
 
            for (auto dop : allDiffs) {
              Value dVal = dop.getInputs()[gradArgPos];
              callerEffectMap[dop].emplace_back(oputils::getEffectOfVal(
                  dVal, MemoryEffects::Write::get(), eff.getResource()));
              callerEffectMap[dop].emplace_back(oputils::getEffectOfVal(
                  dVal, MemoryEffects::Read::get(), eff.getResource()));
            }
          }
        }
 
        if (skipMergeEntry)
          continue;
 
        SmallVector<AutoDiffOp> prunedSources =
            batchutils::pruneGradDefs(key, allDiffs);
 
        SmallVector<AutoDiffOp> legalMerge = batchutils::pruneMemoryEffects(
            symbolTable, key, prunedSources, callerEffectMap, innerEffectCache);
 
        // go ahead and actually do the merge now
        {
 
          SmallVector<enzyme::AutoDiffOp> &allOps = legalMerge;
          int64_t width = allOps.size();
 
          if (width < 2)
            continue;
 
          auto firstDiffOp = allOps.front();
          IRRewriter::InsertionGuard insertGuard(builder);
          builder.setInsertionPoint(firstDiffOp);
          auto loc = firstDiffOp->getLoc();
          auto context = builder.getContext();
 
          // Prepare args for merged operation
 
          SmallVector<mlir::Value> in_args;
          SmallVector<ActivityAttr, 2> inActivityAttrs;
          SmallVector<ActivityAttr, 2> retActivityAttrs;
          SmallVector<mlir::Type, 2> out_ty;
 
          // fill in_args using inputs
          size_t call_idx = 0;
          for (auto [idx, act] : llvm::enumerate(key.inActivity)) {
            auto iattr = ActivityAttr::get(context, act);
            inActivityAttrs.push_back(iattr);
            in_args.push_back(key.inputs[idx]);
            call_idx++;
 
            if (act == Activity::enzyme_dup ||
                act == Activity::enzyme_dupnoneed) {
 
              SmallVector<mlir::Value> derivList;
              for (auto uop : allOps) {
                derivList.push_back(uop.getInputs()[call_idx]);
              }
 
              mlir::Value b_din = getConcatValue(builder, loc, derivList);
              in_args.push_back(b_din);
              call_idx++;
            }
          }
 
          // Skip if function is non-differentiable
          if (call_idx == firstDiffOp.getInputs().size()) {
            continue;
          }
 
          // fill in_args using d<out>, fill out_ty using out
          size_t out_idx = 0;
          for (auto ract : key.retActivity) {
            auto iattr = ActivityAttr::get(context, ract);
            retActivityAttrs.push_back(iattr);
 
            // no effect on out or d<out>
            if (ract == Activity::enzyme_constnoneed ||
                ract == Activity::enzyme_dupnoneed) {
              continue;
            }
 
            // handle d<out>
            if (ract == Activity::enzyme_active ||
                ract == Activity::enzyme_activenoneed) {
              SmallVector<mlir::Value> derivList;
              for (auto uop : allOps) {
                derivList.push_back(uop.getInputs()[call_idx]);
              }
 
              Value batch_dout = getConcatValue(builder, loc, derivList);
              in_args.push_back(batch_dout);
              call_idx++;
            }
 
            // handle out
            if (ract == Activity::enzyme_active ||
                ract == Activity::enzyme_const ||
                ract == Activity::enzyme_dup) {
              Value out = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(out.getType());
              ++out_idx;
            }
          }
 
          // fill out_ty using d<in>
          for (auto act : key.inActivity) {
            if (act == Activity::enzyme_active) {
              Value din = firstDiffOp.getOutputs()[out_idx];
              out_ty.push_back(getConcatType(din, width));
              ++out_idx;
            }
          }
 
          ArrayAttr newInActivity = ArrayAttr::get(
              context, llvm::ArrayRef<Attribute>(inActivityAttrs.begin(),
                                                 inActivityAttrs.end()));
 
          ArrayAttr newRetActivity = ArrayAttr::get(
              context, llvm::ArrayRef<Attribute>(retActivityAttrs.begin(),
                                                 retActivityAttrs.end()));
 
          IntegerAttr newWidthAttr =
              IntegerAttr::get(firstDiffOp.getWidthAttr().getType(), width);
 
          auto newDiffOp = AutoDiffOp::create(
              builder, loc, out_ty, firstDiffOp.getFnAttr(), in_args,
              newInActivity, newRetActivity, newWidthAttr,
              firstDiffOp.getStrongZeroAttr(), firstDiffOp.getAtomicAddAttr());
 
          // Map old uses to new uses
          out_idx = 0;
          for (auto ract : key.retActivity) {
            if (ract == Activity::enzyme_active ||
                ract == Activity::enzyme_const ||
                ract == Activity::enzyme_dup) {
              Value new_out = newDiffOp.getOutputs()[out_idx];
              for (auto dop : allOps) {
                dop.getOutputs()[out_idx].replaceAllUsesWith(new_out);
              }
              ++out_idx;
            }
          }
 
          for (auto act : key.inActivity) {
            if (act == Activity::enzyme_active) {
              Value batch_din = newDiffOp.getOutputs()[out_idx];
              for (auto [dop_idx, dop] : llvm::enumerate(allOps)) {
                Value old_din = dop.getOutputs()[out_idx];
                auto dinTy = old_din.getType();
                auto new_din =
                    getExtractValue(builder, loc, dinTy, batch_din, dop_idx);
 
                old_din.replaceAllUsesWith(new_din);
              }
            }
          }
          // erase all old ops
          for (auto dop : allOps) {
            dop->erase();
          }
        }
      }
    }); // block walker
  }
};
 
} // end anonymous namespace
 
void BatchDiffPass::runOnOperation() {
  SymbolTableCollection symbolTable;
  symbolTable.getSymbolTable(getOperation());
  getOperation()->walk([&](FunctionOpInterface op) {
    mergeFwddiffCalls(symbolTable, op);
    mergeRevdiffCalls(symbolTable, op);
  });
}