CoreDialectsAutoDiffImplementations.h

Go to the documentation of this file.

//===- CoreDialectsAutoDiffImplementation.h - Impl registrations -* C++ -*-===//
//
// 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 contains context registration facilities for external model
// implementations of the automatic differentiation interface for upstream MLIR
// dialects.
//
//===----------------------------------------------------------------------===//
 
#ifndef ENZYMEMLIR_CORE_IMPL_H_
#define ENZYMEMLIR_CORE_IMPL_H_
 
#include "Interfaces/AutoDiffOpInterface.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Support/LogicalResult.h"
 
#include "llvm/ADT/DenseSet.h"
 
namespace mlir {
class DialectRegistry;
class Operation;
class OpBuilder;
class RegionSuccessor;
 
namespace enzyme {
class MGradientUtils;
class MGradientUtilsReverse;
 
namespace detail {
// Non-template implementation of
// AutoDiffUsingControlFlow::createForwardModeTangent.
 
LogicalResult controlFlowForwardHandler(Operation *op, OpBuilder &builder,
                                        MGradientUtils *gutils);
 
LogicalResult controlFlowForwardHandler(
    Operation *op, OpBuilder &builder, MGradientUtils *gutils,
    const llvm::SmallDenseSet<unsigned> &operandPositionsToShadow,
    const llvm::SmallDenseSet<unsigned> &resultPositionsToShadow);
 
// Implements forward-mode differentiation of branching operations.
// Assumes that successive shadows are legal
void branchingForwardHandler(Operation *op, OpBuilder &builder,
                             MGradientUtils *gutils);
 
// Implements forward-mode differentiation of region-terminator operations.
// Assumes that successive shadows are legal
void regionTerminatorForwardHandler(Operation *op, OpBuilder &builder,
                                    MGradientUtils *gutils);
 
// Implements reverse-mode differentiation of return operations.
void returnReverseHandler(Operation *op, OpBuilder &builder,
                          MGradientUtilsReverse *gutils);
 
// Implements forward-mode differentiation of read-only (including read-none)
// operations which do not perform computation
LogicalResult memoryIdentityForwardHandler(Operation *op, OpBuilder &builder,
                                           MGradientUtils *gutils,
                                           ArrayRef<int> storedVals);
 
// Implements shadow initialization differentiation of allocation
LogicalResult allocationForwardHandler(Operation *op, OpBuilder &builder,
                                       MGradientUtils *gutils, bool zero);
 
// Implements the forward autodiff interface for operations whose derivatives
// are can be inferred by analyzing their control flow and differentiating the
// nested operations.
template <typename OpTy>
class AutoDiffUsingControlFlow
    : public AutoDiffOpInterface::ExternalModel<AutoDiffUsingControlFlow<OpTy>,
                                                OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
    return controlFlowForwardHandler(op, builder, gutils);
  }
};
 
// Implements the forward autodiff interface for operations whose derivatives
// are can be inferred by analyzing their branching properties.
template <typename OpTy>
class AutoDiffUsingBranch
    : public AutoDiffOpInterface::ExternalModel<AutoDiffUsingBranch<OpTy>,
                                                OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
    branchingForwardHandler(op, builder, gutils);
    return success();
  }
};
 
// Implements the forward autodiff interface for operations whose derivatives
// are can be inferred by analyzing their region terminator properties.
template <typename OpTy>
class AutoDiffUsingRegionTerminator
    : public AutoDiffOpInterface::ExternalModel<
          AutoDiffUsingRegionTerminator<OpTy>, OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
    regionTerminatorForwardHandler(op, builder, gutils);
    return success();
  }
};
 
template <typename OpTy>
class NoopRevAutoDiffInterface
    : public ReverseAutoDiffOpInterface::ExternalModel<
          NoopRevAutoDiffInterface<OpTy>, OpTy> {
public:
  LogicalResult createReverseModeAdjoint(Operation *op, OpBuilder &builder,
                                         MGradientUtilsReverse *gutils,
                                         SmallVector<Value> caches) const {
    return success();
  }
 
  SmallVector<Value> cacheValues(Operation *op,
                                 MGradientUtilsReverse *gutils) const {
    return SmallVector<Value>();
  }
 
  LogicalResult createShadowValues(Operation *op, OpBuilder &builder,
                                   MGradientUtilsReverse *gutils) const {
    return success();
  }
};
 
template <typename OpTy>
class ReturnRevAutoDiffInterface
    : public ReverseAutoDiffOpInterface::ExternalModel<
          ReturnRevAutoDiffInterface<OpTy>, OpTy> {
public:
  LogicalResult createReverseModeAdjoint(Operation *op, OpBuilder &builder,
                                         MGradientUtilsReverse *gutils,
                                         SmallVector<Value> caches) const {
    returnReverseHandler(op, builder, gutils);
    return success();
  }
 
  SmallVector<Value> cacheValues(Operation *op,
                                 MGradientUtilsReverse *gutils) const {
    return SmallVector<Value>();
  }
 
  LogicalResult createShadowValues(Operation *op, OpBuilder &builder,
                                   MGradientUtilsReverse *gutils) const {
    return success();
  }
};
 
// Implements the forward autodiff interface for operations which are
// read only and identity like (aka not computing sin of mem read).
template <typename OpTy, int... storedvals>
class AutoDiffUsingMemoryIdentity
    : public AutoDiffOpInterface::ExternalModel<
          AutoDiffUsingMemoryIdentity<OpTy, storedvals...>, OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
 
    return memoryIdentityForwardHandler(
        op, builder, gutils, std::initializer_list<int>{storedvals...});
  }
};
 
// Implements the forward autodiff interface for operations which are
// allocation like
template <typename OpTy>
class AutoDiffUsingAllocationFwd : public AutoDiffOpInterface::ExternalModel<
                                       AutoDiffUsingAllocationFwd<OpTy>, OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *op, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
 
    return allocationForwardHandler(op, builder, gutils, /*zero*/ false);
  }
};
 
// Implements the reverse autodiff interface for operations which are
// allocation like
template <typename OpTy>
class AutoDiffUsingAllocationRev
    : public ReverseAutoDiffOpInterface::ExternalModel<
          AutoDiffUsingAllocationRev<OpTy>, OpTy> {
public:
  LogicalResult createReverseModeAdjoint(Operation *op, OpBuilder &builder,
                                         MGradientUtilsReverse *gutils,
                                         SmallVector<Value> caches) const {
    return success();
  }
 
  SmallVector<Value> cacheValues(Operation *op,
                                 MGradientUtilsReverse *gutils) const {
    return SmallVector<Value>();
  }
 
  LogicalResult createShadowValues(Operation *op, OpBuilder &builder,
                                   MGradientUtilsReverse *gutils) const {
    (void)allocationForwardHandler(op, builder, (MGradientUtils *)gutils,
                                   /*zero*/ true);
    return success();
  }
};
 
// Differentiating a call is the same work whichever dialect spelled it:
// find the callee, ask Logic for its derivative, and call that. What differs
// is how the call names its callee and how a call to the derivative is
// written, and both are already asked through interfaces.
 
LogicalResult callForwardHandler(Operation *orig, OpBuilder &builder,
                                 MGradientUtils *gutils);
 
LogicalResult callReverseHandler(Operation *orig, OpBuilder &builder,
                                 MGradientUtilsReverse *gutils,
                                 SmallVector<Value> caches);
 
SmallVector<Value> callCacheValues(Operation *orig,
                                   MGradientUtilsReverse *gutils);
 
} // namespace detail
 
template <typename OpTy>
class AutoDiffCallFwd
    : public AutoDiffOpInterface::ExternalModel<AutoDiffCallFwd<OpTy>, OpTy> {
public:
  LogicalResult createForwardModeTangent(Operation *orig, OpBuilder &builder,
                                         MGradientUtils *gutils) const {
    return detail::callForwardHandler(orig, builder, gutils);
  }
};
 
template <typename OpTy>
class AutoDiffCallRev
    : public ReverseAutoDiffOpInterface::ExternalModel<AutoDiffCallRev<OpTy>,
                                                       OpTy> {
public:
  LogicalResult createReverseModeAdjoint(Operation *orig, OpBuilder &builder,
                                         MGradientUtilsReverse *gutils,
                                         SmallVector<Value> caches) const {
    return detail::callReverseHandler(orig, builder, gutils, caches);
  }
 
  SmallVector<Value> cacheValues(Operation *orig,
                                 MGradientUtilsReverse *gutils) const {
    return detail::callCacheValues(orig, gutils);
  }
 
  LogicalResult createShadowValues(Operation *op, OpBuilder &builder,
                                   MGradientUtilsReverse *gutils) const {
    return success();
  }
};
 
// Registers the generic call handlers for the given op; the tablegen-emitted
// registerInterfaces calls this for each `def : CallOp<...>`.
template <typename OpTy>
void registerAutoDiffUsingCallInterface(MLIRContext &context) {
  OpTy::template attachInterface<AutoDiffCallFwd<OpTy>>(context);
  OpTy::template attachInterface<AutoDiffCallRev<OpTy>>(context);
}
 
// Registers AutoDiffUsingControlFlow for the given op.
template <typename OpTy>
void registerAutoDiffUsingControlFlowInterface(MLIRContext &context) {
  OpTy::template attachInterface<detail::AutoDiffUsingControlFlow<OpTy>>(
      context);
}
// Registers AutoDiffUsingBranch for the given op.
template <typename OpTy>
void registerAutoDiffUsingBranchInterface(MLIRContext &context) {
  OpTy::template attachInterface<detail::AutoDiffUsingBranch<OpTy>>(context);
  OpTy::template attachInterface<detail::NoopRevAutoDiffInterface<OpTy>>(
      context);
}
// Registers AutoDiffUsingRegionTerminator for the given op.
template <typename OpTy>
void registerAutoDiffUsingRegionTerminatorInterface(MLIRContext &context) {
  OpTy::template attachInterface<detail::AutoDiffUsingRegionTerminator<OpTy>>(
      context);
  OpTy::template attachInterface<detail::NoopRevAutoDiffInterface<OpTy>>(
      context);
}
// Registers registerAutoDiffUsingReturnInterface for the given op.
template <typename OpTy>
void registerAutoDiffUsingReturnInterface(MLIRContext &context) {
  OpTy::template attachInterface<detail::AutoDiffUsingRegionTerminator<OpTy>>(
      context);
  OpTy::template attachInterface<detail::ReturnRevAutoDiffInterface<OpTy>>(
      context);
}
// Registers AutoDiffUsingMemoryIdentity for the given op.
template <typename OpTy, int... storedvals>
void registerAutoDiffUsingMemoryIdentityInterface(MLIRContext &context) {
  OpTy::template attachInterface<
      detail::AutoDiffUsingMemoryIdentity<OpTy, storedvals...>>(context);
}
// Registers AutoDiffUsingAllocation for the given op.
template <typename OpTy>
void registerAutoDiffUsingAllocationInterface(MLIRContext &context) {
  OpTy::template attachInterface<detail::AutoDiffUsingAllocationFwd<OpTy>>(
      context);
  OpTy::template attachInterface<detail::AutoDiffUsingAllocationRev<OpTy>>(
      context);
}
 
// Interface registration hooks for individual upstream dialects.
void registerAffineDialectAutoDiffInterface(DialectRegistry &registry);
void registerArithDialectAutoDiffInterface(DialectRegistry &registry);
void registerBuiltinDialectAutoDiffInterface(DialectRegistry &registry);
void registerLLVMDialectAutoDiffInterface(DialectRegistry &registry);
void registerLLVMExtDialectAutoDiffInterface(DialectRegistry &registry);
void registerNVVMDialectAutoDiffInterface(DialectRegistry &registry);
void registerMemRefDialectAutoDiffInterface(DialectRegistry &registry);
void registerComplexDialectAutoDiffInterface(DialectRegistry &registry);
void registerSCFDialectAutoDiffInterface(DialectRegistry &registry);
void registerCFDialectAutoDiffInterface(DialectRegistry &registry);
void registerLinalgDialectAutoDiffInterface(DialectRegistry &registry);
void registerMathDialectAutoDiffInterface(DialectRegistry &registry);
void registerFuncDialectAutoDiffInterface(DialectRegistry &registry);
void registerTensorDialectAutoDiffInterface(DialectRegistry &registry);
void registerGPUDialectAutoDiffInterface(DialectRegistry &registry);
void registerEnzymeDialectAutoDiffInterface(DialectRegistry &registry);
 
void registerCoreDialectAutodiffInterfaces(DialectRegistry &registry);
 
mlir::TypedAttr getConstantAttr(mlir::Type type, llvm::StringRef value);
 
bool allResultsActive(Operation *op, const MGradientUtils *gutils);
 
/// Give an operation built for a derivative the fast-math flags the LLVM path
/// gives everything it builds (Utils.cpp's getFast()). A gradient is only as
/// fast as the arithmetic it is allowed to reassociate, and one built without
/// the flags its primal was compiled with leaves work on the table: much of
/// what the reverse pass computes is a product with a seed, and only under
/// fast-math does a zero seed fold away the expression it feeds.
/// An op with no fast-math flags to give is left alone.
void setDerivativeFastMath(Operation *op);
 
inline Value setDerivativeFastMath(Value v) {
  if (v)
    setDerivativeFastMath(v.getDefiningOp());
  return v;
}
 
template <typename OpTy>
std::enable_if_t<std::is_base_of<OpState, OpTy>::value, OpTy>
setDerivativeFastMath(OpTy op) {
  setDerivativeFastMath(op.getOperation());
  return op;
}
} // namespace enzyme
} // namespace mlir
 
#endif