diff --git a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp index f9510c055dd..bdb569aca0e 100644 --- a/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp +++ b/enzyme/Enzyme/MLIR/Analysis/AliasAnalysis.cpp @@ -38,6 +38,9 @@ // The LLVM dialect is only used for attribute names. #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +// TODO: remove this once aliasing interface is factored out. +#include "mlir/Dialect/MemRef/IR/MemRef.h" + using namespace mlir; using namespace mlir::dataflow; @@ -725,14 +728,28 @@ void enzyme::AliasAnalysis::setToEntryState(AliasClassLattice *lattice) { } } +/// Returns `true` if the alias transfer function of the operation is fully +/// described by its memory effects. +// +// We could have an operation that has side effects and loads a pointer from +// another pointer, but also has another result that aliases the operand, which +// would need additional processing. +// +// TODO: turn this into an interface. +static bool isAliasTransferFullyDescribedByMemoryEffects(Operation *op) { + return isa(op); +} + void enzyme::AliasAnalysis::transfer( Operation *op, ArrayRef effects, ArrayRef operands, ArrayRef results) { + bool globalRead = false; for (const auto &effect : effects) { + // If the effect is global read, record that. Value value = effect.getValue(); if (!value) { - // TODO: we can't assume anything about entry states + globalRead |= isa(effect.getEffect()); continue; } @@ -771,23 +788,32 @@ void enzyme::AliasAnalysis::transfer( } } - // TODO(zinenko): this is sketchy. We could have an operation that has side - // effects and loads a pointer from another pointer, but also has another - // result that aliases the operand. So returning here is premature. - if (!effects.empty()) + // If there was a global read effect, the operation may be reading from any + // pointer so we cannot say what the results are pointing to. Can safely exit + // here because all results are now in the fixpoint state. + if (globalRead) { + for (auto *resultLattice : results) + propagateIfChanged(resultLattice, resultLattice->markUnknown()); + return; + } + + // If it was enough to reason about effects, exit here. + if (!effects.empty() && isAliasTransferFullyDescribedByMemoryEffects(op)) return; - // For operations that don't touch memory, conservatively assume all results - // alias all operands. - for (auto *resultLattice : results) { + // Conservatively assume all results alias all operands. + for (AliasClassLattice *resultLattice : results) { // TODO: Setting this flag to true will assume non-pointers don't alias, // which is not true in general but can result in improved analysis speed. // We need a type analysis for full correctness. - const bool prune_non_pointers = false; - if (!prune_non_pointers || - isPointerLike(resultLattice->getPoint().getType())) - for (const auto *operandLattice : operands) - join(resultLattice, *operandLattice); + constexpr bool pruneNonPointers = false; + if (pruneNonPointers && !isPointerLike(resultLattice->getPoint().getType())) + continue; + + ChangeResult r = ChangeResult::NoChange; + for (const AliasClassLattice *operandLattice : operands) + r |= resultLattice->join(*operandLattice); + propagateIfChanged(resultLattice, r); } }