ActivityAnnotations.h

Go to the documentation of this file.

#ifndef ENZYME_MLIR_ANALYSIS_ACTIVITYANNOTATIONS_H
#define ENZYME_MLIR_ANALYSIS_ACTIVITYANNOTATIONS_H
 
#include "DataFlowAliasAnalysis.h"
#include "DataFlowLattice.h"
#include "Dialect/Ops.h"
 
#include "mlir/Analysis/DataFlow/DenseAnalysis.h"
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
#include "mlir/Analysis/DataFlowFramework.h"
 
namespace mlir {
class FunctionOpInterface;
 
namespace enzyme {
 
using ValueOriginSet = SetLattice<OriginAttr>;
 
//===----------------------------------------------------------------------===//
// ForwardOriginsLattice
//===----------------------------------------------------------------------===//
 
// TODO: specialize this to only arguments
class ForwardOriginsLattice : public SparseSetLattice<OriginAttr> {
public:
  using SparseSetLattice::SparseSetLattice;
 
  static ForwardOriginsLattice single(Value point, OriginAttr value) {
    return ForwardOriginsLattice(point, SetLattice<OriginAttr>(value));
  }
 
  void print(raw_ostream &os) const override;
 
  ChangeResult join(const AbstractSparseLattice &other) override;
 
  const DenseSet<OriginAttr> &getOrigins() const {
    return elements.getElements();
  }
 
  const SetLattice<OriginAttr> &getOriginsObject() const { return elements; }
};
 
class BackwardOriginsLattice : public SparseSetLattice<OriginAttr> {
public:
  using SparseSetLattice::SparseSetLattice;
 
  static BackwardOriginsLattice single(Value point, OriginAttr value) {
    return BackwardOriginsLattice(point, SetLattice<OriginAttr>(value));
  }
 
  void print(raw_ostream &os) const override;
 
  ChangeResult meet(const AbstractSparseLattice &other) override {
    // MLIR framework again misusing terminology
    const auto *otherValueOrigins =
        static_cast<const BackwardOriginsLattice *>(&other);
    return elements.join(otherValueOrigins->elements);
  }
 
  const DenseSet<OriginAttr> &getOrigins() const {
    return elements.getElements();
  }
 
  const SetLattice<OriginAttr> &getOriginsObject() const { return elements; }
};
 
class ForwardActivityAnnotationAnalysis
    : public dataflow::SparseForwardDataFlowAnalysis<ForwardOriginsLattice> {
public:
  ForwardActivityAnnotationAnalysis(DataFlowSolver &solver)
      : SparseForwardDataFlowAnalysis(solver) {
    assert(!solver.getConfig().isInterprocedural());
  }
 
  void setToEntryState(ForwardOriginsLattice *lattice) override;
 
  LogicalResult
  visitOperation(Operation *op,
                 ArrayRef<const ForwardOriginsLattice *> operands,
                 ArrayRef<ForwardOriginsLattice *> results) override;
 
  void visitExternalCall(CallOpInterface call,
                         ArrayRef<const ForwardOriginsLattice *> operands,
                         ArrayRef<ForwardOriginsLattice *> results) override;
 
private:
  void processMemoryRead(Operation *op, Value address,
                         ArrayRef<ForwardOriginsLattice *> results);
 
  void markResultsUnknown(ArrayRef<ForwardOriginsLattice *> results);
 
  void
  processCallToSummarizedFunc(CallOpInterface call,
                              ArrayRef<ValueOriginSet> summary,
                              ArrayRef<const ForwardOriginsLattice *> operands,
                              ArrayRef<ForwardOriginsLattice *> results);
};
 
class BackwardActivityAnnotationAnalysis
    : public dataflow::SparseBackwardDataFlowAnalysis<BackwardOriginsLattice> {
public:
  BackwardActivityAnnotationAnalysis(DataFlowSolver &solver,
                                     SymbolTableCollection &symbolTable)
      : SparseBackwardDataFlowAnalysis(solver, symbolTable) {
    assert(!solver.getConfig().isInterprocedural());
  }
 
  void visitBranchOperand(OpOperand &operand) override {}
 
  void visitCallOperand(OpOperand &operand) override {}
 
  void
  visitNonControlFlowArguments(RegionSuccessor &successor,
                               ArrayRef<BlockArgument> arguments) override {}
 
  void setToExitState(BackwardOriginsLattice *lattice) override;
 
  LogicalResult
  visitOperation(Operation *op, ArrayRef<BackwardOriginsLattice *> operands,
                 ArrayRef<const BackwardOriginsLattice *> results) override;
 
  void
  visitExternalCall(CallOpInterface call,
                    ArrayRef<BackwardOriginsLattice *> operands,
                    ArrayRef<const BackwardOriginsLattice *> results) override;
 
private:
  void
  processCallToSummarizedFunc(CallOpInterface call,
                              ArrayRef<ValueOriginSet> summary,
                              ArrayRef<BackwardOriginsLattice *> operands,
                              ArrayRef<const BackwardOriginsLattice *> results);
 
  void markOperandsUnknown(ArrayRef<BackwardOriginsLattice *> operands);
};
 
//===----------------------------------------------------------------------===//
// ForwardOriginsMap
//===----------------------------------------------------------------------===//
 
class ForwardOriginsMap : public MapOfSetsLattice<DistinctAttr, OriginAttr> {
public:
  using MapOfSetsLattice::MapOfSetsLattice;
 
  void print(raw_ostream &os) const override;
 
  ChangeResult markAllOriginsUnknown() { return markAllUnknown(); }
 
  const ValueOriginSet &getOrigins(DistinctAttr id) const { return lookup(id); }
};
 
class BackwardOriginsMap : public MapOfSetsLattice<DistinctAttr, OriginAttr> {
public:
  using MapOfSetsLattice::MapOfSetsLattice;
 
  void print(raw_ostream &os) const override;
 
  ChangeResult markAllOriginsUnknown() { return markAllUnknown(); }
 
  const ValueOriginSet &getOrigins(DistinctAttr id) const { return lookup(id); }
 
  ChangeResult meet(const AbstractDenseLattice &other) override {
    return join(other);
  }
};
 
//===----------------------------------------------------------------------===//
// DenseActivityAnnotationAnalysis
//===----------------------------------------------------------------------===//
 
class DenseActivityAnnotationAnalysis
    : public dataflow::DenseForwardDataFlowAnalysis<ForwardOriginsMap> {
public:
  using DenseForwardDataFlowAnalysis::DenseForwardDataFlowAnalysis;
 
  void setToEntryState(ForwardOriginsMap *lattice) override;
 
  LogicalResult visitOperation(Operation *op, const ForwardOriginsMap &before,
                               ForwardOriginsMap *after) override;
 
  void visitCallControlFlowTransfer(CallOpInterface call,
                                    dataflow::CallControlFlowAction action,
                                    const ForwardOriginsMap &before,
                                    ForwardOriginsMap *after) override;
 
private:
  void processCallToSummarizedFunc(
      CallOpInterface call,
      const DenseMap<DistinctAttr, ValueOriginSet> &summary,
      const ForwardOriginsMap &before, ForwardOriginsMap *after);
 
  void processCopy(Operation *op, Value copySource, Value copyDest,
                   const ForwardOriginsMap &before, ForwardOriginsMap *after);
 
  OriginalClasses originalClasses;
};
 
class DenseBackwardActivityAnnotationAnalysis
    : public dataflow::DenseBackwardDataFlowAnalysis<BackwardOriginsMap> {
public:
  using DenseBackwardDataFlowAnalysis::DenseBackwardDataFlowAnalysis;
 
  LogicalResult visitOperation(Operation *op, const BackwardOriginsMap &after,
                               BackwardOriginsMap *before) override;
 
  void visitCallControlFlowTransfer(CallOpInterface call,
                                    dataflow::CallControlFlowAction action,
                                    const BackwardOriginsMap &after,
                                    BackwardOriginsMap *before) override;
 
  void setToExitState(BackwardOriginsMap *lattice) override;
 
private:
  void processCallToSummarizedFunc(
      CallOpInterface call,
      const DenseMap<DistinctAttr, ValueOriginSet> &summary,
      const BackwardOriginsMap &after, BackwardOriginsMap *before);
 
  void processCopy(Operation *op, Value copySource, Value copyDest,
                   const BackwardOriginsMap &after, BackwardOriginsMap *before);
};
 
class ActivityPrinterConfig {
public:
  ActivityPrinterConfig() = default;
 
  /// Whether to use the data-flow based algorithm or the classic activity
  /// analysis.
  bool dataflow = true;
  /// Use function summaries
  bool relative = true;
  /// Output extra information for debugging
  bool verbose = false;
  /// Annotate the IR with activity information for every operation. Currently
  /// only supports the LLVM dialect.
  bool annotate = false;
  /// Infer the starting argument state from an __enzyme_autodiff call.
  bool inferFromAutodiff = false;
};
 
void runActivityAnnotations(
    FunctionOpInterface callee, ArrayRef<enzyme::Activity> argActivities = {},
    const ActivityPrinterConfig &config = ActivityPrinterConfig());
 
} // namespace enzyme
} // namespace mlir
 
#endif // ENZYME_MLIR_ANALYSIS_ACTIVITYANNOTATIONS_H