ActivityAnnotations.cpp

Go to the documentation of this file.

#include "ActivityAnnotations.h"
#include "DataFlowAliasAnalysis.h"
#include "DataFlowLattice.h"
#include "Dialect/Dialect.h"
#include "Dialect/Ops.h"
#include "Interfaces/AutoDiffOpInterface.h"
 
#include "mlir/Analysis/DataFlow/DenseAnalysis.h"
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
#include "mlir/Analysis/DataFlow/Utils.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/raw_ostream.h"
 
// TODO: Remove dependency on dialects in favour of differential dependency
// interface
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 
using namespace mlir;
 
static bool isPossiblyActive(Type type) {
  if (auto tensorType = dyn_cast<TensorType>(type)) {
    return isPossiblyActive(tensorType.getElementType());
  }
  return isa<FloatType, ComplexType>(type);
}
 
void enzyme::ForwardOriginsLattice::print(raw_ostream &os) const {
  os << serialize(getAnchor().getContext());
}
 
void enzyme::BackwardOriginsLattice::print(raw_ostream &os) const {
  os << serialize(getAnchor().getContext());
}
 
ChangeResult
enzyme::ForwardOriginsLattice::join(const AbstractSparseLattice &other) {
  const auto *otherValueOrigins =
      static_cast<const ForwardOriginsLattice *>(&other);
  return elements.join(otherValueOrigins->elements);
}
 
void enzyme::ForwardActivityAnnotationAnalysis::setToEntryState(
    ForwardOriginsLattice *lattice) {
  auto arg = dyn_cast<BlockArgument>(lattice->getAnchor());
  if (!arg) {
    assert(lattice->isUndefined());
    return;
  }
  if (!isPossiblyActive(arg.getType())) {
    return;
  }
 
  if (auto funcOp =
          dyn_cast<FunctionOpInterface>(arg.getOwner()->getParentOp())) {
    auto origin = ArgumentOriginAttr::get(FlatSymbolRefAttr::get(funcOp),
                                          arg.getArgNumber());
    return propagateIfChanged(lattice,
                              lattice->join(ForwardOriginsLattice::single(
                                  lattice->getAnchor(), origin)));
  }
}
 
void enzyme::ForwardActivityAnnotationAnalysis::markResultsUnknown(
    ArrayRef<ForwardOriginsLattice *> results) {
  for (ForwardOriginsLattice *result : results) {
    propagateIfChanged(result, result->markUnknown());
  }
}
 
LogicalResult enzyme::ForwardActivityAnnotationAnalysis::visitOperation(
    Operation *op, ArrayRef<const ForwardOriginsLattice *> operands,
    ArrayRef<ForwardOriginsLattice *> results) {
  auto iface = dyn_cast<enzyme::ActivityOpInterface>(op);
  if (iface && iface.isInactive())
    return success();
 
  auto memory = dyn_cast<MemoryEffectOpInterface>(op);
  if (!memory) {
    if (!isMemoryEffectFree(op))
      markResultsUnknown(results);
    return success();
  }
 
  SmallVector<MemoryEffects::EffectInstance> effects;
  memory.getEffects(effects);
  for (const auto &effect : effects) {
    if (!isa<MemoryEffects::Read>(effect.getEffect()))
      continue;
 
    Value value = effect.getValue();
    if (!value) {
      markResultsUnknown(results);
      continue;
    }
    processMemoryRead(op, value, results);
  }
 
  // Propagate sparse differential dependencies, excluding any inactive
  // arguments.
  for (auto &&[i, operand] : llvm::enumerate(operands)) {
    if (iface && iface.isArgInactive(i))
      continue;
 
    for (ForwardOriginsLattice *result : results) {
      join(result, *operand);
    }
  }
  return success();
}
 
void enzyme::ForwardActivityAnnotationAnalysis::processMemoryRead(
    Operation *op, Value address, ArrayRef<ForwardOriginsLattice *> results) {
  ProgramPoint *point = getProgramPointAfter(op);
  auto *srcClasses = getOrCreateFor<AliasClassLattice>(point, address);
  auto *originsMap = getOrCreateFor<ForwardOriginsMap>(point, point);
  if (srcClasses->isUndefined())
    return;
  if (srcClasses->isUnknown())
    return markResultsUnknown(results);
 
  // Look up the alias class and see what its origins are, then propagate
  // those origins to the read results.
  for (DistinctAttr srcClass : srcClasses->getAliasClasses()) {
    for (ForwardOriginsLattice *result : results) {
      if (isPossiblyActive(result->getAnchor().getType())) {
        propagateIfChanged(result,
                           result->merge(originsMap->getOrigins(srcClass)));
      }
    }
  }
}
 
void deserializeReturnOrigins(ArrayAttr returnOrigins,
                              SmallVectorImpl<enzyme::ValueOriginSet> &out) {
  for (auto &&[resultIdx, argOrigins] : llvm::enumerate(returnOrigins)) {
    enzyme::ValueOriginSet origins;
    if (auto strAttr = dyn_cast<StringAttr>(argOrigins)) {
      if (strAttr.getValue() == "<unknown>") {
        (void)origins.markUnknown();
      } else {
        // Leave origins undefined
      }
    } else {
      for (enzyme::ArgumentOriginAttr originAttr :
           cast<ArrayAttr>(argOrigins)
               .getAsRange<enzyme::ArgumentOriginAttr>()) {
        (void)origins.insert({originAttr});
      }
    }
 
    out.push_back(origins);
  }
}
 
void enzyme::ForwardActivityAnnotationAnalysis::visitExternalCall(
    CallOpInterface call, ArrayRef<const ForwardOriginsLattice *> operands,
    ArrayRef<ForwardOriginsLattice *> results) {
  auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
  auto markAllResultsUnknown = [&]() {
    for (ForwardOriginsLattice *result : results) {
      propagateIfChanged(result, result->markUnknown());
    }
  };
  if (!symbol)
    return markAllResultsUnknown();
 
  if (auto callee = SymbolTable::lookupNearestSymbolFrom<FunctionOpInterface>(
          call, symbol.getLeafReference())) {
    if (auto returnOriginsAttr = callee->getAttrOfType<ArrayAttr>(
            EnzymeDialect::getSparseActivityAnnotationAttrName())) {
      SmallVector<ValueOriginSet> returnOrigins;
      deserializeReturnOrigins(returnOriginsAttr, returnOrigins);
      return processCallToSummarizedFunc(call, returnOrigins, operands,
                                         results);
    }
  }
 
  // In the absence of a summary attribute, assume all results differentially
  // depend on all operands
  for (ForwardOriginsLattice *result : results)
    for (const ForwardOriginsLattice *operand : operands)
      join(result, *operand);
}
 
/// Visit everything transitively pointed-to by any pointer in start.
void enzyme::traversePointsToSets(const enzyme::AliasClassSet &start,
                                  const enzyme::PointsToSets &pointsToSets,
                                  function_ref<void(DistinctAttr)> visit) {
  using enzyme::AliasClassSet;
  // The points-to graph is a graph, not a tree: without remembering what has
  // been visited, a cycle spins this loop forever and converging chains
  // revisit their shared tails once per path.
  DenseSet<DistinctAttr> visited;
  AliasClassSet current(start);
  while (!current.isUndefined()) {
    AliasClassSet next;
 
    assert(!current.isUnknown() && "Unhandled traversal of unknown");
    for (DistinctAttr currentClass : current.getElements()) {
      if (!visited.insert(currentClass).second)
        continue;
      visit(currentClass);
      (void)next.join(pointsToSets.getPointsTo(currentClass));
    }
    if (next.isUndefined())
      break;
    std::swap(current, next);
  }
}
 
void enzyme::ForwardActivityAnnotationAnalysis::processCallToSummarizedFunc(
    CallOpInterface call, ArrayRef<ValueOriginSet> summary,
    ArrayRef<const ForwardOriginsLattice *> operands,
    ArrayRef<ForwardOriginsLattice *> results) {
  for (const auto &[result, returnOrigin] : llvm::zip(results, summary)) {
    // Convert the origins relative to the callee to relative to the caller
    ValueOriginSet callerOrigins;
    if (returnOrigin.isUndefined())
      continue;
 
    if (returnOrigin.isUnknown()) {
      (void)callerOrigins.markUnknown();
    } else {
      ProgramPoint *point = getProgramPointAfter(call);
      auto *denseOrigins = getOrCreateFor<ForwardOriginsMap>(point, point);
      auto *pointsTo = getOrCreateFor<PointsToSets>(point, point);
      (void)returnOrigin.foreachElement(
          [&](OriginAttr calleeOrigin, ValueOriginSet::State state) {
            assert(state == ValueOriginSet::State::Defined &&
                   "undefined and unknown must have been handled above");
            auto calleeArgOrigin = cast<ArgumentOriginAttr>(calleeOrigin);
            // If the caller is a pointer, need to join what it points to
            const ForwardOriginsLattice *operandOrigins =
                operands[calleeArgOrigin.getArgNumber()];
            auto *callerAliasClass = getOrCreateFor<AliasClassLattice>(
                getProgramPointAfter(call), operandOrigins->getAnchor());
            traversePointsToSets(callerAliasClass->getAliasClassesObject(),
                                 *pointsTo, [&](DistinctAttr aliasClass) {
                                   (void)callerOrigins.join(
                                       denseOrigins->getOrigins(aliasClass));
                                 });
            return callerOrigins.join(operandOrigins->getOriginsObject());
          });
    }
    propagateIfChanged(result, result->merge(callerOrigins));
  }
}
 
void enzyme::BackwardActivityAnnotationAnalysis::setToExitState(
    BackwardOriginsLattice *lattice) {}
 
void enzyme::BackwardActivityAnnotationAnalysis::markOperandsUnknown(
    ArrayRef<BackwardOriginsLattice *> operands) {
  for (BackwardOriginsLattice *operand : operands) {
    propagateIfChanged(operand, operand->markUnknown());
  }
}
 
LogicalResult enzyme::BackwardActivityAnnotationAnalysis::visitOperation(
    Operation *op, ArrayRef<BackwardOriginsLattice *> operands,
    ArrayRef<const BackwardOriginsLattice *> results) {
  auto iface = dyn_cast<enzyme::ActivityOpInterface>(op);
  if (iface && iface.isInactive())
    return success();
 
  auto memory = dyn_cast<MemoryEffectOpInterface>(op);
  if (!memory) {
    if (!isMemoryEffectFree(op))
      markOperandsUnknown(operands);
    return success();
  }
 
  SmallVector<MemoryEffects::EffectInstance> effects;
  memory.getEffects(effects);
  for (const auto &effect : effects) {
    if (!isa<MemoryEffects::Read>(effect.getEffect()))
      continue;
 
    Value value = effect.getValue();
    if (!value) {
      markOperandsUnknown(operands);
      continue;
    }
 
    auto *srcClasses =
        getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), value);
    auto *originsMap =
        getOrCreate<BackwardOriginsMap>(getProgramPointBefore(op));
 
    ChangeResult changed = ChangeResult::NoChange;
    for (const BackwardOriginsLattice *result : results)
      changed |= originsMap->insert(srcClasses->getAliasClassesObject(),
                                    result->getOriginsObject());
    propagateIfChanged(originsMap, changed);
  }
 
  // Propagate sparse differential dependencies, excluding any inactive
  // arguments.
  for (auto &&[i, operand] : llvm::enumerate(operands)) {
    if (iface && iface.isArgInactive(i))
      continue;
 
    for (const BackwardOriginsLattice *result : results)
      meet(operand, *result);
  }
  return success();
}
 
void enzyme::BackwardActivityAnnotationAnalysis::visitExternalCall(
    CallOpInterface call, ArrayRef<BackwardOriginsLattice *> operands,
    ArrayRef<const BackwardOriginsLattice *> results) {
  auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
  if (!symbol)
    return markOperandsUnknown(operands);
 
  if (auto callee = SymbolTable::lookupNearestSymbolFrom<FunctionOpInterface>(
          call, symbol.getLeafReference())) {
    if (auto returnOriginsAttr = callee->getAttrOfType<ArrayAttr>(
            EnzymeDialect::getSparseActivityAnnotationAttrName())) {
      SmallVector<ValueOriginSet> returnOrigins;
      deserializeReturnOrigins(returnOriginsAttr, returnOrigins);
      return processCallToSummarizedFunc(call, returnOrigins, operands,
                                         results);
    }
  }
 
  // In the absence of a summary attribute, assume all results differentially
  // depend on all operands
  for (BackwardOriginsLattice *operand : operands)
    for (const BackwardOriginsLattice *result : results)
      meet(operand, *result);
}
 
void enzyme::BackwardActivityAnnotationAnalysis::processCallToSummarizedFunc(
    CallOpInterface call, ArrayRef<ValueOriginSet> summary,
    ArrayRef<BackwardOriginsLattice *> operands,
    ArrayRef<const BackwardOriginsLattice *> results) {
  // collect the result origins, propagate them to the operands.
  for (const auto &[result, calleeOrigins] : llvm::zip(results, summary)) {
    ValueOriginSet resultOrigins = result->getOriginsObject();
    if (calleeOrigins.isUndefined())
      continue;
    if (calleeOrigins.isUnknown())
      (void)resultOrigins.markUnknown();
    else {
      (void)calleeOrigins.foreachElement(
          [&](OriginAttr calleeOrigin, ValueOriginSet::State state) {
            auto calleeArgOrigin = cast<ArgumentOriginAttr>(calleeOrigin);
            BackwardOriginsLattice *operand =
                operands[calleeArgOrigin.getArgNumber()];
            propagateIfChanged(operand, operand->merge(resultOrigins));
            return ChangeResult::NoChange;
          });
    }
  }
}
 
template <typename KeyT, typename ElementT>
void printMapOfSetsLattice(
    const DenseMap<KeyT, enzyme::SetLattice<ElementT>> map, raw_ostream &os) {
  if (map.empty()) {
    os << "<empty>\n";
    return;
  }
  for (const auto &[aliasClass, origins] : map) {
    os << "  " << aliasClass << " originates from {";
    if (origins.isUnknown()) {
      os << "<unknown>";
    } else if (origins.isUndefined()) {
      os << "<undefined>";
    } else {
      llvm::interleaveComma(origins.getElements(), os);
    }
    os << "}\n";
  }
}
 
void enzyme::ForwardOriginsMap::print(raw_ostream &os) const {
  printMapOfSetsLattice(this->map, os);
}
 
void enzyme::BackwardOriginsMap::print(raw_ostream &os) const {
  printMapOfSetsLattice(this->map, os);
}
 
void enzyme::DenseActivityAnnotationAnalysis::setToEntryState(
    ForwardOriginsMap *lattice) {
  auto point = dyn_cast<ProgramPoint *>(lattice->getAnchor());
  auto *block = point->getBlock();
  if (!(block && isa<FunctionOpInterface>(block->getParentOp())))
    return;
 
  auto funcOp = cast<FunctionOpInterface>(block->getParentOp());
  ChangeResult changed = ChangeResult::NoChange;
  for (BlockArgument arg : funcOp.getArguments()) {
    auto *argClass = getOrCreateFor<AliasClassLattice>(point, arg);
    auto origin = ArgumentOriginAttr::get(FlatSymbolRefAttr::get(funcOp),
                                          arg.getArgNumber());
    changed |= lattice->insert(argClass->getAliasClassesObject(),
                               ValueOriginSet(origin));
  }
  propagateIfChanged(lattice, changed);
}
 
std::optional<Value> getStored(Operation *op);
std::optional<Value> getCopySource(Operation *op);
 
LogicalResult enzyme::DenseActivityAnnotationAnalysis::visitOperation(
    Operation *op, const ForwardOriginsMap &before, ForwardOriginsMap *after) {
  join(after, before);
 
  auto iface = dyn_cast<enzyme::ActivityOpInterface>(op);
  if (iface && iface.isInactive())
    return success();
 
  auto memory = dyn_cast<MemoryEffectOpInterface>(op);
  if (!memory) {
    if (!isMemoryEffectFree(op))
      propagateIfChanged(after, after->markAllOriginsUnknown());
    return success();
  }
 
  SmallVector<MemoryEffects::EffectInstance> effects;
  memory.getEffects(effects);
  for (const auto &effect : effects) {
    Value value = effect.getValue();
    if (!value) {
      propagateIfChanged(after, after->markAllOriginsUnknown());
      return success();
    }
 
    if (isa<MemoryEffects::Read>(effect.getEffect())) {
      // TODO: Really need that memory interface
      if (op->getNumResults() != 1)
        continue;
      Value readDest = op->getResult(0);
 
      auto *destClasses =
          getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), readDest);
      if (destClasses->isUndefined())
        // Not a pointer, so the sparse analysis will handle this.
        continue;
 
      auto *srcClasses =
          getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), value);
      if (srcClasses->isUnknown()) {
        propagateIfChanged(after,
                           after->insert(destClasses->getAliasClassesObject(),
                                         ValueOriginSet::getUnknown()));
        continue;
      }
 
      ChangeResult changed = ChangeResult::NoChange;
      for (DistinctAttr srcClass : srcClasses->getAliasClasses()) {
        changed |= after->insert(destClasses->getAliasClassesObject(),
                                 before.getOrigins(srcClass));
      }
      propagateIfChanged(after, changed);
    } else if (isa<MemoryEffects::Write>(effect.getEffect())) {
      if (std::optional<Value> stored = getStored(op)) {
        if (!isPossiblyActive(stored->getType())) {
          continue;
        }
        auto *origins = getOrCreateFor<ForwardOriginsLattice>(
            getProgramPointAfter(op), *stored);
        auto *dest =
            getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), value);
        propagateIfChanged(after, after->insert(dest->getAliasClassesObject(),
                                                origins->getOriginsObject()));
      } else if (std::optional<Value> copySource = getCopySource(op)) {
        processCopy(op, *copySource, value, before, after);
      } else {
        propagateIfChanged(after, after->markAllOriginsUnknown());
      }
    }
  }
  return success();
}
 
void enzyme::DenseActivityAnnotationAnalysis::processCopy(
    Operation *op, Value copySource, Value copyDest,
    const ForwardOriginsMap &before, ForwardOriginsMap *after) {
  auto *src =
      getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), copySource);
  ValueOriginSet srcOrigins;
  if (src->isUndefined())
    return;
  if (src->isUnknown())
    (void)srcOrigins.markUnknown();
 
  for (DistinctAttr srcClass : src->getAliasClasses())
    (void)srcOrigins.join(before.getOrigins(srcClass));
 
  auto *dest =
      getOrCreateFor<AliasClassLattice>(getProgramPointAfter(op), copyDest);
  propagateIfChanged(after,
                     after->insert(dest->getAliasClassesObject(), srcOrigins));
}
 
// TODO: rename from pointsto
static void deserializePointsTo(
    ArrayAttr summaryAttr,
    DenseMap<DistinctAttr, enzyme::ValueOriginSet> &summaryMap) {
  // TODO: investigate better encodings for the value origin summary
  for (auto pair : summaryAttr.getAsRange<ArrayAttr>()) {
    assert(pair.size() == 2 &&
           "Expected summary to be in [[key, value]] format");
    auto pointer = cast<DistinctAttr>(pair[0]);
    auto pointsToSet = enzyme::ValueOriginSet::getUndefined();
    if (auto strAttr = dyn_cast<StringAttr>(pair[1])) {
      if (strAttr.getValue() == enzyme::unknownSetString) {
        (void)pointsToSet.markUnknown();
      } else {
        assert(strAttr.getValue() == enzyme::undefinedSetString &&
               "unrecognized points-to destination");
      }
    } else {
      auto pointsTo = cast<ArrayAttr>(pair[1]).getAsRange<enzyme::OriginAttr>();
      (void)pointsToSet.insert(
          DenseSet<enzyme::OriginAttr>(pointsTo.begin(), pointsTo.end()));
    }
 
    summaryMap.insert({pointer, pointsToSet});
  }
}
 
void enzyme::DenseActivityAnnotationAnalysis::visitCallControlFlowTransfer(
    CallOpInterface call, dataflow::CallControlFlowAction action,
    const ForwardOriginsMap &before, ForwardOriginsMap *after) {
  join(after, before);
  if (action == dataflow::CallControlFlowAction::ExternalCallee) {
    auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
    if (!symbol)
      return propagateIfChanged(after, after->markAllOriginsUnknown());
 
    if (auto callee = SymbolTable::lookupNearestSymbolFrom<FunctionOpInterface>(
            call, symbol.getLeafReference())) {
      if (auto summaryAttr = callee->getAttrOfType<ArrayAttr>(
              EnzymeDialect::getDenseActivityAnnotationAttrName())) {
        DenseMap<DistinctAttr, ValueOriginSet> summary;
        deserializePointsTo(summaryAttr, summary);
        return processCallToSummarizedFunc(call, summary, before, after);
      }
    }
  }
}
 
void enzyme::DenseActivityAnnotationAnalysis::processCallToSummarizedFunc(
    CallOpInterface call, const DenseMap<DistinctAttr, ValueOriginSet> &summary,
    const ForwardOriginsMap &before, ForwardOriginsMap *after) {
  ChangeResult changed = ChangeResult::NoChange;
  ProgramPoint *point = getProgramPointAfter(call);
  // Unify the value origin summary with the actual lattices of function
  // arguments
  // Collect the origins of the function arguments, then collect the alias
  // classes of the destinations
  auto *p2sets = getOrCreateFor<PointsToSets>(point, point);
  SmallVector<ValueOriginSet> argumentOrigins;
  SmallVector<AliasClassSet> argumentClasses;
  for (auto &&[i, argOperand] : llvm::enumerate(call.getArgOperands())) {
    // Value origin might be sparse, might be dense
    ValueOriginSet argOrigins;
    auto *argClasses = getOrCreateFor<AliasClassLattice>(point, argOperand);
    if (argClasses->isUndefined()) {
      // Not a pointer, use the sparse lattice state
      auto *sparseOrigins =
          getOrCreateFor<ForwardOriginsLattice>(point, argOperand);
      (void)argOrigins.join(sparseOrigins->getOriginsObject());
    } else {
      // Unify all the origins
      // Since we're not keeping track of argument depth, we need to union the
      // arg origins with everything it points to.
      traversePointsToSets(argClasses->getAliasClassesObject(), *p2sets,
                           [&](DistinctAttr aliasClass) {
                             (void)argOrigins.join(
                                 before.getOrigins(aliasClass));
                           });
    }
    argumentClasses.push_back(argClasses->getAliasClassesObject());
    argumentOrigins.push_back(argOrigins);
  }
 
  for (const auto &[destClass, sourceOrigins] : summary) {
    ValueOriginSet callerOrigins;
    // A callee that could not say where this memory came from says nothing
    // narrower to the caller either; only a set that names its origins has
    // arguments to translate.
    if (sourceOrigins.isUndefined())
      continue;
    if (sourceOrigins.isUnknown()) {
      (void)callerOrigins.markUnknown();
    } else {
      for (Attribute sourceOrigin : sourceOrigins.getElements()) {
        unsigned argNumber =
            cast<ArgumentOriginAttr>(sourceOrigin).getArgNumber();
        (void)callerOrigins.join(argumentOrigins[argNumber]);
      }
    }
 
    AliasClassSet callerDestClasses;
    if (auto pseudoClass = dyn_cast_if_present<PseudoAliasClassAttr>(
            destClass.getReferencedAttr())) {
      // Traverse the points-to sets.
      AliasClassSet current = argumentClasses[pseudoClass.getArgNumber()];
      unsigned depth = pseudoClass.getDepth();
      while (depth > 0) {
        AliasClassSet next;
        if (current.isUndefined()) {
          // Activity annotations requires converged pointer info. If the alias
          // class is undefined, this could be because (1) the points-to info
          // hasn't _yet_ been computed (in which case we bail out here
          // expecting to be called again with more complete info), or (2) if
          // the points-to info has converged, this signifies reading from
          // uninitialized memory. This is UB, so we assume it never happens.
          return;
        }
        for (DistinctAttr currentClass : current.getElements())
          (void)next.join(p2sets->getPointsTo(currentClass));
        std::swap(current, next);
        depth--;
      }
 
      (void)callerDestClasses.join(current);
    } else {
      (void)callerDestClasses.insert({destClass});
    }
    changed |= after->insert(callerDestClasses, callerOrigins);
  }
  propagateIfChanged(after, changed);
}
 
void enzyme::DenseBackwardActivityAnnotationAnalysis::
    visitCallControlFlowTransfer(CallOpInterface call,
                                 dataflow::CallControlFlowAction action,
                                 const BackwardOriginsMap &after,
                                 BackwardOriginsMap *before) {
  meet(before, after);
  if (action == dataflow::CallControlFlowAction::ExternalCallee) {
    auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
    if (!symbol)
      return propagateIfChanged(before, before->markAllOriginsUnknown());
 
    if (auto callee = SymbolTable::lookupNearestSymbolFrom<FunctionOpInterface>(
            call, symbol.getLeafReference())) {
      if (auto summaryAttr = callee->getAttrOfType<ArrayAttr>(
              EnzymeDialect::getDenseActivityAnnotationAttrName())) {
        DenseMap<DistinctAttr, ValueOriginSet> summary;
        deserializePointsTo(summaryAttr, summary);
        return processCallToSummarizedFunc(call, summary, after, before);
      }
    }
  }
}
 
void enzyme::DenseBackwardActivityAnnotationAnalysis::setToExitState(
    BackwardOriginsMap *lattice) {
  auto point = dyn_cast<ProgramPoint *>(lattice->getAnchor());
  auto *block = point->getBlock();
  if (!block)
    return;
 
  auto funcOp = cast<FunctionOpInterface>(block->getParentOp());
  ChangeResult changed = ChangeResult::NoChange;
  for (BlockArgument arg : funcOp.getArguments()) {
    auto *pointsToSets = getOrCreateFor<PointsToSets>(
        point, getProgramPointAfter(block->getTerminator()));
    auto *argClass = getOrCreateFor<AliasClassLattice>(point, arg);
    auto origin = ArgumentOriginAttr::get(FlatSymbolRefAttr::get(funcOp),
                                          arg.getArgNumber());
 
    // Everything that a pointer argument may point to originates from that
    // pointer argument.
    traversePointsToSets(argClass->getAliasClassesObject(), *pointsToSets,
                         [&](DistinctAttr currentClass) {
                           changed |=
                               lattice->insert(AliasClassSet(currentClass),
                                               ValueOriginSet(origin));
                         });
  }
  propagateIfChanged(lattice, changed);
}
 
LogicalResult enzyme::DenseBackwardActivityAnnotationAnalysis::visitOperation(
    Operation *op, const BackwardOriginsMap &after,
    BackwardOriginsMap *before) {
  meet(before, after);
 
  auto activityIface = dyn_cast<enzyme::ActivityOpInterface>(op);
  if (activityIface && activityIface.isInactive())
    return success();
 
  auto memory = dyn_cast<MemoryEffectOpInterface>(op);
  if (!memory) {
    if (!isMemoryEffectFree(op))
      propagateIfChanged(before, before->markAllOriginsUnknown());
    return success();
  }
 
  SmallVector<MemoryEffects::EffectInstance> effects;
  memory.getEffects(effects);
  for (const auto &effect : effects) {
    if (!isa<MemoryEffects::Write>(effect.getEffect()))
      continue;
 
    Value value = effect.getValue();
    if (!value) {
      propagateIfChanged(before, before->markAllOriginsUnknown());
      return success();
    }
 
    if (std::optional<Value> stored = getStored(op)) {
      ProgramPoint *point = getProgramPointBefore(op);
      auto *addressClasses = getOrCreateFor<AliasClassLattice>(point, value);
      auto *storedClasses = getOrCreateFor<AliasClassLattice>(point, *stored);
 
      if (storedClasses->isUndefined()) {
        // Not a pointer being stored, do a sparse update
        auto *storedOrigins = getOrCreate<BackwardOriginsLattice>(*stored);
        propagateIfChanged(
            storedOrigins,
            addressClasses->getAliasClassesObject().foreachElement(
                [&](DistinctAttr alloc, AliasClassSet::State state) {
                  if (state == AliasClassSet::State::Undefined) {
                    return ChangeResult::NoChange;
                  }
                  if (state == AliasClassSet::State::Unknown) {
                    return storedOrigins->markUnknown();
                  }
                  return storedOrigins->merge(after.getOrigins(alloc));
                }));
      } else if (storedClasses->isUnknown()) {
        propagateIfChanged(before, before->markAllOriginsUnknown());
      } else {
        // Capturing stores are handled via the points-to relationship in
        // setToExitState.
      }
    } else if (std::optional<Value> copySource = getCopySource(op)) {
      processCopy(op, *copySource, value, after, before);
    }
  }
  return success();
}
 
void enzyme::DenseBackwardActivityAnnotationAnalysis::
    processCallToSummarizedFunc(
        CallOpInterface call,
        const DenseMap<DistinctAttr, ValueOriginSet> &summary,
        const BackwardOriginsMap &after, BackwardOriginsMap *before) {
  ChangeResult changed = ChangeResult::NoChange;
  ProgramPoint *pointBefore = getProgramPointBefore(call);
  ProgramPoint *pointAfter = getProgramPointAfter(call);
  // Unify the value origin summary with the actual lattices of function
  // arguments
  auto *p2sets = getOrCreateFor<PointsToSets>(pointBefore, pointAfter);
  SmallVector<AliasClassSet> argumentClasses;
  for (Value argOperand : call.getArgOperands()) {
    auto *argClasses =
        getOrCreateFor<AliasClassLattice>(pointBefore, argOperand);
    argumentClasses.push_back(argClasses->getAliasClassesObject());
  }
 
  for (const auto &[destClass, sourceOrigins] : summary) {
    // Get the destination origins
    ValueOriginSet destOrigins;
    if (auto pseudoClass = dyn_cast_if_present<PseudoAliasClassAttr>(
            destClass.getReferencedAttr())) {
      traversePointsToSets(argumentClasses[pseudoClass.getArgNumber()], *p2sets,
                           [&](DistinctAttr aliasClass) {
                             (void)destOrigins.join(
                                 after.getOrigins(aliasClass));
                           });
    }
 
    if (destOrigins.isUndefined())
      continue;
 
    // A callee that could not say which of its arguments this memory flows back
    // to leaves every one of the caller's classes a candidate.
    if (sourceOrigins.isUndefined())
      continue;
    if (sourceOrigins.isUnknown()) {
      changed |= before->markAllOriginsUnknown();
      continue;
    }
 
    // Get the source alias classes
    AliasClassSet callerSourceClasses;
    for (Attribute sourceOrigin : sourceOrigins.getElements()) {
      unsigned argNumber =
          cast<ArgumentOriginAttr>(sourceOrigin).getArgNumber();
 
      if (argumentClasses[argNumber].isUndefined()) {
        // Not a pointer, do a sparse update
        auto *backwardLattice = getOrCreate<BackwardOriginsLattice>(
            call.getArgOperands()[argNumber]);
        if (destOrigins.isUnknown()) {
          propagateIfChanged(backwardLattice, backwardLattice->markUnknown());
          continue;
        }
        propagateIfChanged(backwardLattice,
                           backwardLattice->insert(destOrigins.getElements()));
      } else {
        traversePointsToSets(argumentClasses[argNumber], *p2sets,
                             [&](DistinctAttr aliasClass) {
                               (void)callerSourceClasses.insert({aliasClass});
                             });
      }
    }
    changed |= before->insert(callerSourceClasses, destOrigins);
  }
  propagateIfChanged(before, changed);
}
 
void enzyme::DenseBackwardActivityAnnotationAnalysis::processCopy(
    Operation *op, Value copySource, Value copyDest,
    const BackwardOriginsMap &after, BackwardOriginsMap *before) {
  ProgramPoint *point = getProgramPointBefore(op);
  auto *dest = getOrCreateFor<AliasClassLattice>(point, copyDest);
  ValueOriginSet destOrigins;
  if (dest->isUndefined())
    return;
  if (dest->isUnknown())
    (void)destOrigins.markUnknown();
 
  for (DistinctAttr destClass : dest->getAliasClasses())
    (void)destOrigins.join(after.getOrigins(destClass));
 
  auto *src = getOrCreateFor<AliasClassLattice>(point, copySource);
  propagateIfChanged(before,
                     before->insert(src->getAliasClassesObject(), destOrigins));
}
 
namespace {
 
// TODO: the alias summary attribute is sufficent to get the correct behaviour
// here, but it would be nice if these were not hardcoded.
void annotateHardcoded(FunctionOpInterface func) {
  if (func.getName() == "lgamma" || func.getName() == "tanh") {
    MLIRContext *ctx = func.getContext();
    SmallVector<Attribute> arr = {StringAttr::get(ctx, "<undefined>")};
    func->setAttr(enzyme::EnzymeDialect::getAliasSummaryAttrName(),
                  ArrayAttr::get(ctx, arr));
  }
}
 
void initializeSparseBackwardActivityAnnotations(FunctionOpInterface func,
                                                 DataFlowSolver &solver) {
  using namespace mlir::enzyme;
 
  for (Operation &op : func.getCallableRegion()->getOps()) {
    if (!mlir::enzyme::isFunctionReturnLike(&op))
      continue;
 
    for (OpOperand &returnOperand : op.getOpOperands()) {
      auto *lattice =
          solver.getOrCreateState<BackwardOriginsLattice>(returnOperand.get());
      auto origin = ReturnOriginAttr::get(FlatSymbolRefAttr::get(func),
                                          returnOperand.getOperandNumber());
      (void)lattice->insert({origin});
    }
  }
}
 
using OriginsPair =
    std::pair<enzyme::ForwardOriginsLattice, enzyme::BackwardOriginsLattice>;
 
/// Once having reached a top-level entry point, go top-down and convert the
/// relative sources/sinks into concrete active/constant results.
///
/// This would ideally be done after lowering to LLVM and during differentiation
/// because it loses context sensitivity, but this is faster to prototype with.
void topDownActivityAnalysis(
    FunctionOpInterface callee, ArrayRef<enzyme::Activity> argActivities,
    ArrayRef<enzyme::Activity> retActivities,
    DenseMap<BlockArgument, OriginsPair> &blockArgOrigins) {
  using namespace mlir::enzyme;
  MLIRContext *ctx = callee.getContext();
  callee->setAttr("enzyme.visited", UnitAttr::get(ctx));
  auto trueAttr = BoolAttr::get(ctx, true);
  auto falseAttr = BoolAttr::get(ctx, false);
 
  auto isOriginActive = [&](OriginAttr origin) {
    if (auto argOriginAttr = dyn_cast<ArgumentOriginAttr>(origin)) {
      return llvm::is_contained({Activity::enzyme_dup,
                                 Activity::enzyme_dupnoneed,
                                 Activity::enzyme_active},
                                argActivities[argOriginAttr.getArgNumber()]);
    }
    auto retOriginAttr = cast<ReturnOriginAttr>(origin);
    return llvm::is_contained({Activity::enzyme_dup, Activity::enzyme_dupnoneed,
                               Activity::enzyme_active},
                              retActivities[retOriginAttr.getReturnNumber()]);
  };
  callee.getFunctionBody().walk([&](Operation *op) {
    if (op->getNumResults() == 0) {
      // Operations that don't return values are definitionally "constant"
      op->setAttr("enzyme.icv", trueAttr);
    } else {
      // Value activity
      if (op->hasAttr("enzyme.constantval")) {
        op->setAttr("enzyme.icv", trueAttr);
      } else if (op->hasAttr("enzyme.activeval")) {
        op->setAttr("enzyme.icv", falseAttr);
      } else {
        auto valueSource = op->getAttrOfType<ArrayAttr>("enzyme.valsrc");
        auto valueSink = op->getAttrOfType<ArrayAttr>("enzyme.valsink");
        if (!(valueSource && valueSink)) {
          llvm::errs() << "[activity] missing attributes for op: " << *op
                       << "\n";
        }
        assert(valueSource && valueSink && "missing attributes for op");
        bool activeSource =
            llvm::any_of(valueSource.getAsRange<OriginAttr>(), isOriginActive);
        bool activeSink =
            llvm::any_of(valueSink.getAsRange<OriginAttr>(), isOriginActive);
        bool activeVal = activeSource && activeSink;
        op->setAttr("enzyme.icv", BoolAttr::get(ctx, !activeVal));
      }
    }
    op->removeAttr("enzyme.constantval");
    op->removeAttr("enzyme.activeval");
    op->removeAttr("enzyme.valsrc");
    op->removeAttr("enzyme.valsink");
 
    // Instruction activity
    if (op->hasAttr("enzyme.constantop")) {
      op->setAttr("enzyme.ici", trueAttr);
    } else if (op->hasAttr("enzyme.activeop")) {
      op->setAttr("enzyme.ici", falseAttr);
    } else {
      bool activeSource = llvm::any_of(
          op->getAttrOfType<ArrayAttr>("enzyme.opsrc").getAsRange<OriginAttr>(),
          isOriginActive);
      bool activeSink =
          llvm::any_of(op->getAttrOfType<ArrayAttr>("enzyme.opsink")
                           .getAsRange<OriginAttr>(),
                       isOriginActive);
      bool activeOp = activeSource && activeSink;
      op->setAttr("enzyme.ici", BoolAttr::get(ctx, !activeOp));
    }
 
    op->removeAttr("enzyme.constantop");
    op->removeAttr("enzyme.activeop");
    op->removeAttr("enzyme.opsrc");
    op->removeAttr("enzyme.opsink");
 
    if (auto callOp = dyn_cast<CallOpInterface>(op)) {
      auto funcOp = cast<FunctionOpInterface>(callOp.resolveCallable());
      if (!funcOp->hasAttr("enzyme.visited")) {
        SmallVector<Activity> callArgActivities, callResActivities;
        for (Value operand : callOp.getArgOperands()) {
          if (auto *definingOp = operand.getDefiningOp()) {
            bool icv =
                definingOp->getAttrOfType<BoolAttr>("enzyme.icv").getValue();
            callArgActivities.push_back(icv ? Activity::enzyme_const
                                            : Activity::enzyme_active);
          } else {
            BlockArgument blockArg = cast<BlockArgument>(operand);
            const OriginsPair &originsPair = blockArgOrigins.at(blockArg);
            const ForwardOriginsLattice &sources = originsPair.first;
            const BackwardOriginsLattice &sinks = originsPair.second;
            bool argActive = false;
            if (sources.isUnknown() || sinks.isUnknown()) {
              argActive = true;
            } else if (sources.isUndefined() || sinks.isUndefined()) {
              argActive = false;
            } else {
              argActive = llvm::any_of(sources.getOrigins(), isOriginActive) &&
                          llvm::any_of(sinks.getOrigins(), isOriginActive);
            }
            callArgActivities.push_back(argActive ? Activity::enzyme_active
                                                  : Activity::enzyme_const);
          }
        }
        if (op->getNumResults() != 0) {
          bool icv = op->getAttrOfType<BoolAttr>("enzyme.icv").getValue();
          callResActivities.push_back(icv ? Activity::enzyme_const
                                          : Activity::enzyme_active);
        }
 
        topDownActivityAnalysis(funcOp, callArgActivities, callResActivities,
                                blockArgOrigins);
      }
    }
  });
}
 
// Convenience for making debugging less verbose
std::string format(Attribute attribute) {
  return llvm::TypeSwitch<Attribute, std::string>(attribute)
      .Case([](DistinctAttr attr) { return format(attr.getReferencedAttr()); })
      .Case([](ArrayAttr attr) {
        std::string res;
        llvm::raw_string_ostream ss(res);
        ss << "[";
        for (unsigned i = 0; i < attr.size(); ++i) {
          ss << format(attr[i]);
          if (i < attr.size() - 1)
            ss << ", ";
        }
        ss << "]";
        return res;
      })
      .Case([](enzyme::PseudoAliasClassAttr attr) {
        return ("pclass<(" + Twine(attr.getArgNumber()) + ", " +
                Twine(attr.getDepth()) + ")>")
            .str();
      })
      .Case([](enzyme::ArgumentOriginAttr attr) {
        return ("arg" + Twine(attr.getArgNumber())).str();
      })
      .Case([](enzyme::ReturnOriginAttr attr) {
        return ("ret" + Twine(attr.getReturnNumber())).str();
      })
      .Default([](Attribute attr) {
        std::string res;
        llvm::raw_string_ostream ss(res);
        ss << attr;
        return res;
      });
}
} // namespace
 
/// Starting from callee, compute a reverse (bottom-up) topological sorting of
/// all functions transitively called from callee.
LogicalResult
enzyme::reverseToposortCallgraph(CallableOpInterface callee,
                                 SymbolTableCollection *symbolTable,
                                 SmallVectorImpl<CallableOpInterface> &sorted) {
  DenseSet<CallableOpInterface> permanent;
  DenseSet<CallableOpInterface> temporary;
  std::function<LogicalResult(CallableOpInterface)> visit =
      [&](CallableOpInterface node) {
        if (permanent.contains(node))
          return success();
        if (temporary.contains(node))
          // Cycle in call graph, toposort not possible
          return failure();
 
        temporary.insert(node);
        auto walkResult = node.walk([&](CallOpInterface call) {
          auto neighbour = cast<CallableOpInterface>(
              call.resolveCallableInTable(symbolTable));
          if (failed(visit(neighbour)))
            return WalkResult::interrupt();
          return WalkResult::advance();
        });
        if (walkResult.wasInterrupted())
          return failure();
 
        temporary.erase(node);
        permanent.insert(node);
        sorted.push_back(node);
        return success();
      };
 
  return visit(callee);
}
 
void enzyme::computeSummaries(
    FunctionOpInterface funcOp, DataFlowSolver &solver,
    enzyme::PointsToSets &p2sets, enzyme::ForwardOriginsMap &forwardOriginsMap,
    enzyme::BackwardOriginsMap &backwardOriginsMap,
    SmallVectorImpl<enzyme::ForwardOriginsLattice> &returnOperandOrigins,
    SmallVectorImpl<enzyme::AliasClassLattice> &returnAliasClasses) {
  assert(returnOperandOrigins.size() == funcOp.getNumResults());
  assert(returnAliasClasses.size() == funcOp.getNumResults());
 
  // Create the overall summary by joining sets at all return sites.
  for (Operation &op : funcOp.getCallableRegion()->getOps()) {
    if (mlir::enzyme::isFunctionReturnLike(&op)) {
      ProgramPoint *point = solver.getProgramPointAfter(&op);
      (void)p2sets.join(*solver.lookupState<enzyme::PointsToSets>(point));
      auto *returnOrigins =
          solver.lookupState<enzyme::ForwardOriginsMap>(point);
      if (returnOrigins)
        (void)forwardOriginsMap.join(*returnOrigins);
 
      for (OpOperand &operand : op.getOpOperands()) {
        (void)returnAliasClasses[operand.getOperandNumber()].join(
            *solver.lookupState<enzyme::AliasClassLattice>(operand.get()));
        (void)returnOperandOrigins[operand.getOperandNumber()].join(
            *solver.lookupState<enzyme::ForwardOriginsLattice>(operand.get()));
      }
    }
  }
 
  (void)backwardOriginsMap.meet(
      *solver.getOrCreateState<enzyme::BackwardOriginsMap>(
          solver.getProgramPointBefore(
              &funcOp.getCallableRegion()->front().front())));
}
 
void enzyme::serializeSummaries(
    Operation *op, const enzyme::PointsToSets &p2sets,
    const enzyme::ForwardOriginsMap &forwardOriginsMap,
    ArrayRef<enzyme::ForwardOriginsLattice> returnOperandOrigins,
    ArrayRef<enzyme::AliasClassLattice> returnAliasClasses) {
  using namespace mlir::enzyme;
  MLIRContext *ctx = op->getContext();
  // Sparse alias annotations
  SmallVector<Attribute> aliasAttributes(returnAliasClasses.size());
  llvm::transform(returnAliasClasses, aliasAttributes.begin(),
                  [&](enzyme::AliasClassLattice lattice) {
                    return lattice.serialize(ctx);
                  });
  op->setAttr(EnzymeDialect::getAliasSummaryAttrName(),
              ArrayAttr::get(ctx, aliasAttributes));
 
  // Points-to-pointer annotations
  StringRef pointerSummaryName = EnzymeDialect::getPointerSummaryAttrName();
  op->setAttr(pointerSummaryName, p2sets.serialize(ctx));
 
  op->setAttr(EnzymeDialect::getDenseActivityAnnotationAttrName(),
              forwardOriginsMap.serialize(ctx));
 
  // Serialize return origins
  SmallVector<Attribute> serializedReturnOperandOrigins(
      returnOperandOrigins.size());
  llvm::transform(returnOperandOrigins, serializedReturnOperandOrigins.begin(),
                  [ctx](enzyme::ForwardOriginsLattice lattice) -> Attribute {
                    return lattice.serialize(ctx);
                  });
  op->setAttr(EnzymeDialect::getSparseActivityAnnotationAttrName(),
              ArrayAttr::get(ctx, serializedReturnOperandOrigins));
}
 
void enzyme::removeSummaries(Operation *op) {
  op->removeAttr(EnzymeDialect::getAliasSummaryAttrName());
  op->removeAttr(EnzymeDialect::getPointerSummaryAttrName());
  op->removeAttr(EnzymeDialect::getDenseActivityAnnotationAttrName());
  op->removeAttr(EnzymeDialect::getSparseActivityAnnotationAttrName());
}
 
LogicalResult
enzyme::runActivityAnnotationsForFunction(FunctionOpInterface funcOp,
                                          DataFlowSolver &solver) {
  using namespace mlir::enzyme;
 
  SymbolTableCollection symbolTable;
  MLIRContext *ctx = funcOp.getContext();
  dataflow::loadBaselineAnalyses(solver);
  solver.load<enzyme::AliasAnalysis>(ctx,
                                     /*relative=*/true);
  solver.load<enzyme::PointsToPointerAnalysis>();
  solver.load<enzyme::ForwardActivityAnnotationAnalysis>();
  solver.load<enzyme::DenseActivityAnnotationAnalysis>();
  solver.load<enzyme::BackwardActivityAnnotationAnalysis>(symbolTable);
  solver.load<enzyme::DenseBackwardActivityAnnotationAnalysis>(symbolTable);
 
  initializeSparseBackwardActivityAnnotations(funcOp, solver);
  return solver.initializeAndRun(funcOp);
}
 
void enzyme::runActivityAnnotations(
    FunctionOpInterface callee, ArrayRef<enzyme::Activity> argActivities,
    const ActivityPrinterConfig &activityConfig) {
  SymbolTableCollection symbolTable;
  SmallVector<CallableOpInterface> sorted;
  (void)reverseToposortCallgraph(callee, &symbolTable, sorted);
  raw_ostream &os = llvm::outs();
 
  // TODO: is there any way of serializing information in a block argument?
  DenseMap<BlockArgument, OriginsPair> blockArgOrigins;
 
  StringRef pointerSummaryName = EnzymeDialect::getPointerSummaryAttrName();
  for (CallableOpInterface node : sorted) {
    annotateHardcoded(cast<FunctionOpInterface>(node.getOperation()));
 
    if (!node.getCallableRegion() || node->hasAttr(pointerSummaryName))
      continue;
    auto funcOp = cast<FunctionOpInterface>(node.getOperation());
    if (activityConfig.verbose) {
      os << "[ata] processing function @" << funcOp.getName() << "\n";
    }
    DataFlowConfig dataFlowConfig;
    dataFlowConfig.setInterprocedural(false);
    DataFlowSolver solver(dataFlowConfig);
    if (failed(runActivityAnnotationsForFunction(funcOp, solver))) {
      assert(false && "dataflow solver failed");
    }
 
    // Create the overall summary by joining sets at all return sites.
    enzyme::PointsToSets p2sets(nullptr);
    enzyme::ForwardOriginsMap forwardOriginsMap(nullptr);
    enzyme::BackwardOriginsMap backwardOriginsMap(nullptr);
    size_t numResults = funcOp.getResultTypes().size();
    SmallVector<enzyme::ForwardOriginsLattice> returnOperandOrigins(
        numResults, ForwardOriginsLattice(nullptr));
    SmallVector<enzyme::AliasClassLattice> returnAliasClasses(
        numResults, AliasClassLattice(nullptr));
    computeSummaries(funcOp, solver, p2sets, forwardOriginsMap,
                     backwardOriginsMap, returnOperandOrigins,
                     returnAliasClasses);
    serializeSummaries(funcOp, p2sets, forwardOriginsMap, returnOperandOrigins,
                       returnAliasClasses);
    Attribute backwardOrigins = backwardOriginsMap.serialize(node.getContext());
 
    if (activityConfig.verbose) {
      os << "[ata] p2p summary:\n";
      if (node->getAttrOfType<ArrayAttr>(pointerSummaryName).size() == 0) {
        os << "     <empty>\n";
      }
      for (ArrayAttr pair : node->getAttrOfType<ArrayAttr>(pointerSummaryName)
                                .getAsRange<ArrayAttr>()) {
        os << "     " << format(pair[0]) << " -> " << format(pair[1]) << "\n";
      }
 
      os << "[ata] forward value origins:\n";
      for (ArrayAttr pair :
           node->getAttrOfType<ArrayAttr>(
                   EnzymeDialect::getDenseActivityAnnotationAttrName())
               .getAsRange<ArrayAttr>()) {
        os << "     " << format(pair[0]) << " originates from "
           << format(pair[1]) << "\n";
      }
 
      os << "[ata] backward value origins:\n";
      for (ArrayAttr pair :
           cast<ArrayAttr>(backwardOrigins).getAsRange<ArrayAttr>()) {
        os << "     " << format(pair[0]) << " goes to " << format(pair[1])
           << "\n";
      }
 
      os << "[ata] return origins: "
         << format(node->getAttr(
                EnzymeDialect::getSparseActivityAnnotationAttrName()))
         << "\n";
    }
 
    auto joinActiveDataState =
        [&](Value value,
            std::pair<ForwardOriginsLattice, BackwardOriginsLattice> &out) {
          auto *sources = solver.getOrCreateState<ForwardOriginsLattice>(value);
          auto *sinks = solver.getOrCreateState<BackwardOriginsLattice>(value);
          (void)out.first.join(*sources);
          (void)out.second.meet(*sinks);
        };
 
    auto joinActivePointerState =
        [&](const AliasClassSet &aliasClasses,
            std::pair<ForwardOriginsLattice, BackwardOriginsLattice> &out) {
          traversePointsToSets(
              aliasClasses, p2sets, [&](DistinctAttr aliasClass) {
                (void)out.first.merge(forwardOriginsMap.getOrigins(aliasClass));
                (void)out.second.merge(
                    backwardOriginsMap.getOrigins(aliasClass));
              });
        };
 
    auto joinActiveValueState =
        [&](Value value,
            std::pair<ForwardOriginsLattice, BackwardOriginsLattice> &out) {
          if (isa<LLVM::LLVMPointerType, MemRefType>(value.getType())) {
            auto *aliasClasses =
                solver.getOrCreateState<AliasClassLattice>(value);
            joinActivePointerState(aliasClasses->getAliasClassesObject(), out);
          } else {
            joinActiveDataState(value, out);
          }
        };
 
    MLIRContext *ctx = node.getContext();
    auto annotateActivity = [&](Operation *op) {
      assert(op->getNumResults() < 2 && op->getNumRegions() == 0 &&
             "annotation only supports the LLVM dialect");
      auto unitAttr = UnitAttr::get(ctx);
      // Check activity of values
      for (OpResult result : op->getResults()) {
        std::pair<ForwardOriginsLattice, BackwardOriginsLattice>
            activityAttributes({result, ValueOriginSet()},
                               {result, ValueOriginSet()});
        joinActiveValueState(result, activityAttributes);
        const auto &sources = activityAttributes.first;
        const auto &sinks = activityAttributes.second;
        // Possible states: if either source or sink is undefined or empty, the
        // value is always constant.
        if (sources.isUnknown() || sinks.isUnknown()) {
          // Always active
          op->setAttr("enzyme.activeval", unitAttr);
        } else if (sources.isUndefined() || sinks.isUndefined()) {
          // Always constant
          op->setAttr("enzyme.constantval", unitAttr);
        } else {
          // Conditionally active depending on the activity of sources and sinks
          op->setAttr("enzyme.valsrc", sources.serialize(ctx));
          op->setAttr("enzyme.valsink", sinks.serialize(ctx));
        }
      }
      // Check activity of operation
      StringRef opSourceAttrName = "enzyme.opsrc";
      StringRef opSinkAttrName = "enzyme.opsink";
      std::pair<ForwardOriginsLattice, BackwardOriginsLattice> opAttributes(
          {nullptr, ValueOriginSet()}, {nullptr, ValueOriginSet()});
      if (isPure(op)) {
        // A pure operation can only propagate data via its results
        for (OpResult result : op->getResults()) {
          joinActiveDataState(result, opAttributes);
        }
      } else {
        // We need a special case because stores of active pointers don't fit
        // the definition but are active instructions
        if (auto storeOp = dyn_cast<LLVM::StoreOp>(op)) {
          auto *storedClass =
              solver.getOrCreateState<AliasClassLattice>(storeOp.getValue());
          joinActivePointerState(storedClass->getAliasClassesObject(),
                                 opAttributes);
        } else if (auto callOp = dyn_cast<CallOpInterface>(op)) {
          // TODO: tricky, requires some thought
          auto callable = cast<CallableOpInterface>(callOp.resolveCallable());
          if (callable->hasAttr(
                  EnzymeDialect::getDenseActivityAnnotationAttrName())) {
            for (Value operand : callOp.getArgOperands())
              joinActiveValueState(operand, opAttributes);
          }
          // We need to
          // determine if the body of the function contains active instructions
        }
 
        // Default: the op is active iff any of its operands or results are
        // active data.
        for (Value operand : op->getOperands())
          joinActiveDataState(operand, opAttributes);
        for (OpResult result : op->getResults())
          joinActiveDataState(result, opAttributes);
      }
 
      const auto &opSources = opAttributes.first;
      const auto &opSinks = opAttributes.second;
      if (opSources.isUnknown() || opSinks.isUnknown()) {
        op->setAttr("enzyme.activeop", unitAttr);
      } else if (opSources.isUndefined() || opSinks.isUndefined()) {
        op->setAttr("enzyme.constantop", unitAttr);
      } else {
        op->setAttr(opSourceAttrName, opAttributes.first.serialize(ctx));
        op->setAttr(opSinkAttrName, opAttributes.second.serialize(ctx));
      }
    };
 
    // We lose the solver state when going top down and I don't know a better
    // way to serialize block argument information.
    node.getCallableRegion()->walk([&](Block *block) {
      for (BlockArgument blockArg : block->getArguments()) {
        OriginsPair blockArgAttributes({blockArg, ValueOriginSet()},
                                       {blockArg, ValueOriginSet()});
        joinActiveValueState(blockArg, blockArgAttributes);
        blockArgOrigins.try_emplace(blockArg, blockArgAttributes);
      }
    });
 
    node.getCallableRegion()->walk([&](Operation *op) {
      if (activityConfig.annotate)
        annotateActivity(op);
      if (activityConfig.verbose) {
        if (op->hasAttr("tag")) {
          for (OpResult result : op->getResults()) {
            std::pair<ForwardOriginsLattice, BackwardOriginsLattice>
                activityAttributes({result, ValueOriginSet()},
                                   {result, ValueOriginSet()});
            joinActiveValueState(result, activityAttributes);
            os << op->getAttr("tag") << "(#" << result.getResultNumber()
               << ")\n"
               << "  sources: "
               << format(activityAttributes.first.serialize(ctx)) << "\n"
               << "  sinks:   "
               << format(activityAttributes.second.serialize(ctx)) << "\n";
          }
        }
      }
    });
  }
 
  if (!argActivities.empty() && activityConfig.annotate) {
    SmallVector<enzyme::Activity> resActivities;
    for (Type resultType : callee.getResultTypes()) {
      resActivities.push_back(isPossiblyActive(resultType)
                                  ? Activity::enzyme_active
                                  : Activity::enzyme_const);
    }
 
    topDownActivityAnalysis(callee, argActivities, resActivities,
                            blockArgOrigins);
  }
}