diff --git a/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp b/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp index 0a5e5e8b431..6fb52f4c6fc 100644 --- a/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp +++ b/enzyme/Enzyme/MLIR/Analysis/ActivityAnalysis.cpp @@ -935,15 +935,17 @@ bool mlir::enzyme::ActivityAnalyzer::isConstantOperation(MTypeResults const &TR, // Even if returning a pointer, this instruction is considered inactive // since the instruction doesn't prop gradients. Thus, so long as we don't // return an object containing a float, this instruction is inactive - // if (!TR.intType(1, I, /*errifNotFound*/ false).isPossibleFloat()) { - // if (EnzymePrintActivity) - // llvm::errs() - // << " constant instruction from known non-float non-writing " - // "instruction " - // << *I << "\n"; - // InsertConstantInstruction(TR, I); - // return true; - // } + if (I->getNumResults() != 0 && llvm::all_of(I->getResults(), [&](Value v) { + return !TR.intType(1, v, /*errIfNotFound*/ false).isPossibleFloat(); + })) { + if (EnzymePrintActivity) + llvm::errs() + << " constant instruction from known non-float non-writing " + "instruction " + << *I << "\n"; + InsertConstantOperation(TR, I); + return true; + } // If all returned values constant otherwise, the operation is inactive if (llvm::all_of(I->getResults(), @@ -1599,20 +1601,17 @@ bool mlir::enzyme::ActivityAnalyzer::isConstantValue(MTypeResults const &TR, return true; } -#if 0 - // This value is certainly a pointer to an integer (and only and integer, not + // This value is certainly a pointer to an integer (and only an integer, not // a pointer or float). Therefore its value is constant - // TODO use typeInfo for more aggressive activity analysis - if (val->getType()->isPointerTy() && - cast(val->getType())->isIntOrIntVectorTy() && - TR.firstPointer(1, val, /*errifnotfound*/ false).isIntegral()) { + auto mutableTy = dyn_cast(Val.getType()); + if (mutableTy && mutableTy.isMutable() && + TR.firstPointer(1, Val, /*errIfNotFound*/ false).isIntegral()) { if (EnzymePrintActivity) llvm::errs() << " Value const as integral pointer" << (int)directions - << " " << *val << "\n"; + << " " << Val << "\n"; InsertConstantValue(TR, Val); return true; } -#endif // TODO: since in MLIR globals are operations like others, this should be // handled via interfaces. diff --git a/enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt b/enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt index 0799cc8985a..f0324448e17 100644 --- a/enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt +++ b/enzyme/Enzyme/MLIR/Interfaces/CMakeLists.txt @@ -19,6 +19,7 @@ add_mlir_library( TransformUtils.cpp EnzymeLogic.cpp EnzymeLogicReverse.cpp + MTypeResults.cpp Utils.cpp DEPENDS MLIROffsetViewInterfaceIncGen diff --git a/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h b/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h index 5f752fdd2e3..a73088ec2ae 100644 --- a/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h +++ b/enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h @@ -4,9 +4,12 @@ #include "mlir/IR/IRMapping.h" #include "mlir/Interfaces/FunctionInterfaces.h" -#include "../../TypeAnalysis/TypeAnalysis.h" +#include "../../TypeAnalysis/TypeTree.h" #include "../../Utils.h" #include +#include +#include +#include namespace mlir { namespace enzyme { @@ -27,21 +30,65 @@ class MTypeAnalysis { } }; +} // namespace enzyme +} // namespace mlir + +template <> struct ConcreteTypeTraits { + static bool isFloatingPoint(mlir::Type T) { return static_cast(T); } + + static void print(llvm::raw_ostream &OS, mlir::Type T) { OS << T; } + + static std::string floatName(mlir::Type T) { + std::string S; + llvm::raw_string_ostream OS(S); + OS << T; + return S; + } + + static bool less(mlir::Type A, mlir::Type B) { + return A.getAsOpaquePointer() < B.getAsOpaquePointer(); + } +}; + +template <> struct TypeTreeTraits { + using Origin = mlir::Operation *; + + // Mirror the defaults of -enzyme-max-type-depth and -enzyme-max-type-offset. + static unsigned maxDepth() { return 6; } + + static int maxOffset() { return 500; } + + template + static void depthLimitExceeded(const TreeT &Tree, Origin Orig, + const llvm::Twine &Detail) { + std::string Msg = ("not handling more than " + llvm::Twine(maxDepth()) + + " pointer lookups deep dt:" + Tree.str() + Detail) + .str(); + if (Orig) + Orig->emitWarning(Msg); + else + llvm::errs() << Msg << "\n"; + } +}; + +namespace mlir { +namespace enzyme { + +using MConcreteType = GenericConcreteType; + +using MTypeTree = GenericTypeTree; + +// The MLIR counterpart of TypeResults. There is no type analysis behind it. class MTypeResults { public: - // TODO - TypeTree getReturnAnalysis() { return TypeTree(); } - TypeTree query(Value) const { return TypeTree(); } - ConcreteType intType(size_t num, Value val, bool errIfNotFound = true, - bool pointerIntSame = false) const { - if (isa(val.getType())) { - return BaseType::Integer; - } - if (errIfNotFound) { - llvm_unreachable("something happened"); - } - return BaseType::Unknown; - } + MTypeTree query(Value val) const; + + /// What the value itself is. + MConcreteType intType(size_t num, Value val, bool errIfNotFound = true, + bool pointerIntSame = false) const; + + MConcreteType firstPointer(size_t num, Value val, bool errIfNotFound = true, + bool pointerIntSame = false) const; }; class MEnzymeLogic { diff --git a/enzyme/Enzyme/MLIR/Interfaces/MTypeResults.cpp b/enzyme/Enzyme/MLIR/Interfaces/MTypeResults.cpp new file mode 100644 index 00000000000..f163c21ac30 --- /dev/null +++ b/enzyme/Enzyme/MLIR/Interfaces/MTypeResults.cpp @@ -0,0 +1,44 @@ +//===- MTypeResults.cpp - Type facts for MLIR activity analysis -----------===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// MTypeResults answers the activity analysis's questions about what a value +// holds. Nothing is inferred: every answer is read off the value's type, which +// is exact where the type says what a value holds and Unknown elsewhere. +// +//===----------------------------------------------------------------------===// + +#include "EnzymeLogic.h" + +using namespace mlir; +using namespace mlir::enzyme; + +MTypeTree MTypeResults::query(Value val) const { + // TODO: cache the tree per type in a DenseMap once + // building it is more than a few isa checks. + MTypeTree tree; + return tree; +} + +MConcreteType MTypeResults::intType(size_t num, Value val, bool errIfNotFound, + bool pointerIntSame) const { + if (isa(val.getType())) + return BaseType::Integer; + if (errIfNotFound) + llvm_unreachable("MTypeResults::intType: could not deduce the type of the " + "value from its MLIR type"); + return BaseType::Unknown; +} + +MConcreteType MTypeResults::firstPointer(size_t num, Value val, + bool errIfNotFound, + bool pointerIntSame) const { + if (errIfNotFound) + llvm_unreachable("MTypeResults::firstPointer: could not deduce the pointee " + "type of the value from its MLIR type"); + return BaseType::Unknown; +} diff --git a/enzyme/Enzyme/TypeAnalysis/ConcreteType.h b/enzyme/Enzyme/TypeAnalysis/ConcreteType.h index 8708fbdfcc9..22323409831 100644 --- a/enzyme/Enzyme/TypeAnalysis/ConcreteType.h +++ b/enzyme/Enzyme/TypeAnalysis/ConcreteType.h @@ -24,6 +24,10 @@ // an the SubType category ``BaseType`` as well as the SubType of float, if // relevant. This also contains several helper utility functions. // +// The class is templated over the handle of the floating point SubType so that +// it is shared between LLVM (``llvm::Type *``) and MLIR (``mlir::Type``). +// Everything specific to the IR is provided by ``ConcreteTypeTraits``. +// //===----------------------------------------------------------------------===// #ifndef ENZYME_TYPE_ANALYSIS_CONCRETE_TYPE_H #define ENZYME_TYPE_ANALYSIS_CONCRETE_TYPE_H 1 @@ -33,63 +37,99 @@ #include "llvm/IR/InstrTypes.h" #include "llvm/IR/Type.h" #include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/raw_ostream.h" #include "BaseType.h" +template struct ConcreteTypeTraits; + +template <> struct ConcreteTypeTraits { + static bool isFloatingPoint(llvm::Type *T) { + return !llvm::isa(T) && T->isFloatingPointTy(); + } + + static void print(llvm::raw_ostream &OS, llvm::Type *T) { OS << *T; } + + static std::string floatName(llvm::Type *T) { + if (T->isHalfTy()) + return "half"; + if (T->isFloatTy()) + return "float"; + if (T->isDoubleTy()) + return "double"; + if (T->isX86_FP80Ty()) + return "fp80"; + if (T->isBFloatTy()) + return "bf16"; + if (T->isFP128Ty()) + return "fp128"; + if (T->isPPC_FP128Ty()) + return "ppc128"; + llvm_unreachable("unknown data SubType"); + } + + static llvm::Type *parseFloat(llvm::StringRef SubName, llvm::LLVMContext &C) { + if (SubName == "half") + return llvm::Type::getHalfTy(C); + if (SubName == "float") + return llvm::Type::getFloatTy(C); + if (SubName == "double") + return llvm::Type::getDoubleTy(C); + if (SubName == "fp80") + return llvm::Type::getX86_FP80Ty(C); + if (SubName == "bf16") + return llvm::Type::getBFloatTy(C); + if (SubName == "fp128") + return llvm::Type::getFP128Ty(C); + if (SubName == "ppc128") + return llvm::Type::getPPC_FP128Ty(C); + llvm_unreachable("unknown data SubType"); + } + + static bool less(llvm::Type *A, llvm::Type *B) { return A < B; } +}; + /// Concrete SubType of a given value. Consists of a category `BaseType` and the /// particular floating point value, if relevant. -class ConcreteType { +template class GenericConcreteType { + using Traits = ConcreteTypeTraits; + public: /// Category of underlying type BaseType SubTypeEnum; - /// Floating point type, if relevant, otherwise nullptr - llvm::Type *SubType; + /// Floating point type, if relevant, otherwise null + SubTypeT SubType; /// Construct a ConcreteType from an existing FloatingPoint Type - ConcreteType(llvm::Type *SubType) + GenericConcreteType(SubTypeT SubType) : SubTypeEnum(BaseType::Float), SubType(SubType) { - assert(SubType != nullptr); - assert(!llvm::isa(SubType)); - if (!SubType->isFloatingPointTy()) { - llvm::errs() << " passing in non FP SubType: " << *SubType << "\n"; + assert(SubType); + if (!Traits::isFloatingPoint(SubType)) { + llvm::errs() << " passing in non FP SubType: "; + Traits::print(llvm::errs(), SubType); + llvm::errs() << "\n"; } - assert(SubType->isFloatingPointTy()); + assert(Traits::isFloatingPoint(SubType)); } /// Construct a non-floating Concrete type from a BaseType - ConcreteType(BaseType SubTypeEnum) - : SubTypeEnum(SubTypeEnum), SubType(nullptr) { + GenericConcreteType(BaseType SubTypeEnum) + : SubTypeEnum(SubTypeEnum), SubType() { assert(SubTypeEnum != BaseType::Float); } /// Construct a ConcreteType from a string /// A Concrete Type's string representation is given by the string of the /// enum If it is a floating point it is given by Float@ - ConcreteType(llvm::StringRef Str, llvm::LLVMContext &C) { + template + GenericConcreteType(llvm::StringRef Str, ContextT &C) { auto Sep = Str.find('@'); if (Sep != llvm::StringRef::npos) { SubTypeEnum = BaseType::Float; assert(Str.substr(0, Sep) == "Float"); - auto SubName = Str.substr(Sep + 1); - if (SubName == "half") { - SubType = llvm::Type::getHalfTy(C); - } else if (SubName == "float") { - SubType = llvm::Type::getFloatTy(C); - } else if (SubName == "double") { - SubType = llvm::Type::getDoubleTy(C); - } else if (SubName == "fp80") { - SubType = llvm::Type::getX86_FP80Ty(C); - } else if (SubName == "bf16") { - SubType = llvm::Type::getBFloatTy(C); - } else if (SubName == "fp128") { - SubType = llvm::Type::getFP128Ty(C); - } else if (SubName == "ppc128") { - SubType = llvm::Type::getPPC_FP128Ty(C); - } else { - llvm_unreachable("unknown data SubType"); - } + SubType = Traits::parseFloat(Str.substr(Sep + 1), C); } else { - SubType = nullptr; + SubType = SubTypeT(); SubTypeEnum = parseBaseType(Str); } } @@ -97,25 +137,8 @@ class ConcreteType { /// Convert the ConcreteType to a string std::string str() const { std::string Result = to_string(SubTypeEnum); - if (SubTypeEnum == BaseType::Float) { - if (SubType->isHalfTy()) { - Result += "@half"; - } else if (SubType->isFloatTy()) { - Result += "@float"; - } else if (SubType->isDoubleTy()) { - Result += "@double"; - } else if (SubType->isX86_FP80Ty()) { - Result += "@fp80"; - } else if (SubType->isBFloatTy()) { - Result += "@bf16"; - } else if (SubType->isFP128Ty()) { - Result += "@fp128"; - } else if (SubType->isPPC_FP128Ty()) { - Result += "@ppc128"; - } else { - llvm_unreachable("unknown data SubType"); - } - } + if (SubTypeEnum == BaseType::Float) + Result += "@" + Traits::floatName(SubType); return Result; } @@ -142,7 +165,7 @@ class ConcreteType { } /// Return the floating point type, if this is a float - llvm::Type *isFloat() const { return SubType; } + SubTypeT isFloat() const { return SubType; } /// Return if this is known to be the BaseType BT /// This cannot be called with BaseType::Float as it lacks information @@ -169,16 +192,16 @@ class ConcreteType { } /// Return if this is known to be the ConcreteType CT - bool operator==(const ConcreteType CT) const { + bool operator==(const GenericConcreteType CT) const { return SubType == CT.SubType && SubTypeEnum == CT.SubTypeEnum; } /// Return if this is known not to be the ConcreteType CT - bool operator!=(const ConcreteType CT) const { return !(*this == CT); } + bool operator!=(const GenericConcreteType CT) const { return !(*this == CT); } /// Set this to the given ConcreteType, returning true if /// this ConcreteType has changed - bool operator=(const ConcreteType CT) { + bool operator=(const GenericConcreteType CT) { bool changed = false; if (SubTypeEnum != CT.SubTypeEnum) changed = true; @@ -193,13 +216,14 @@ class ConcreteType { /// this ConcreteType has changed bool operator=(const BaseType BT) { assert(BT != BaseType::Float); - return ConcreteType::operator=(ConcreteType(BT)); + return GenericConcreteType::operator=(GenericConcreteType(BT)); } /// Set this to the logical or of itself and CT, returning whether this value /// changed Setting `PointerIntSame` considers pointers and integers as /// equivalent If this is an illegal operation, `LegalOr` will be set to false - bool checkedOrIn(const ConcreteType CT, bool PointerIntSame, bool &LegalOr) { + bool checkedOrIn(const GenericConcreteType CT, bool PointerIntSame, + bool &LegalOr) { if (SubTypeEnum == BaseType::Anything) { return false; } @@ -236,7 +260,7 @@ class ConcreteType { /// Set this to the logical or of itself and CT, returning whether this value /// changed Setting `PointerIntSame` considers pointers and integers as /// equivalent This function will error if doing an illegal Operation - bool orIn(const ConcreteType CT, bool PointerIntSame) { + bool orIn(const GenericConcreteType CT, bool PointerIntSame) { bool Legal = true; bool Result = checkedOrIn(CT, PointerIntSame, Legal); if (!Legal) { @@ -251,14 +275,14 @@ class ConcreteType { /// Set this to the logical or of itself and CT, returning whether this value /// changed This assumes that pointers and integers are distinct This function /// will error if doing an illegal Operation - bool operator|=(const ConcreteType CT) { + bool operator|=(const GenericConcreteType CT) { return orIn(CT, /*pointerIntSame*/ false); } /// Set this to the logical and of itself and CT, returning whether this value /// changed If this and CT are incompatible, the result will be /// BaseType::Unknown - bool andIn(const ConcreteType CT) { + bool andIn(const GenericConcreteType CT) { if (SubTypeEnum == BaseType::Anything) { return *this = CT; } @@ -284,10 +308,10 @@ class ConcreteType { /// Set this to the logical and of itself and CT, returning whether this value /// changed If this and CT are incompatible, the result will be /// BaseType::Unknown - bool operator&=(const ConcreteType CT) { return andIn(CT); } + bool operator&=(const GenericConcreteType CT) { return andIn(CT); } /// Keep only mappings where the type is not an `Anything` - ConcreteType PurgeAnything() const { + GenericConcreteType PurgeAnything() const { if (SubTypeEnum == BaseType::Anything) return BaseType::Unknown; return *this; @@ -296,7 +320,7 @@ class ConcreteType { /// Set this to the logical `binop` of itself and RHS, using the Binop Op, /// returning true if this was changed. /// This function will error on an invalid type combination - bool binopIn(bool &Legal, const ConcreteType RHS, + bool binopIn(bool &Legal, const GenericConcreteType RHS, llvm::BinaryOperator::BinaryOps Op) { bool Changed = false; using namespace llvm; @@ -316,7 +340,7 @@ class ConcreteType { RHS.SubTypeEnum == BaseType::Integer || RHS.SubTypeEnum == BaseType::Unknown)))) { SubTypeEnum = BaseType::Unknown; - SubType = nullptr; + SubType = SubTypeT(); Changed = true; return Changed; } @@ -503,16 +527,21 @@ class ConcreteType { } /// Compare concrete types for use in map's - bool operator<(const ConcreteType dt) const { + bool operator<(const GenericConcreteType dt) const { if (SubTypeEnum == dt.SubTypeEnum) { - return SubType < dt.SubType; + return Traits::less(SubType, dt.SubType); } else { return SubTypeEnum < dt.SubTypeEnum; } } }; +using ConcreteType = GenericConcreteType; + // Convert ConcreteType to string -static inline std::string to_string(const ConcreteType dt) { return dt.str(); } +template +static inline std::string to_string(const GenericConcreteType dt) { + return dt.str(); +} #endif diff --git a/enzyme/Enzyme/TypeAnalysis/TypeTree.h b/enzyme/Enzyme/TypeAnalysis/TypeTree.h index bb5a7484025..b869fe07db0 100644 --- a/enzyme/Enzyme/TypeAnalysis/TypeTree.h +++ b/enzyme/Enzyme/TypeAnalysis/TypeTree.h @@ -61,7 +61,38 @@ static inline std::string to_string(const std::vector x) { return out; } -class TypeTree; +template struct TypeTreeTraits; + +template <> struct TypeTreeTraits { + using Origin = llvm::Instruction *; + + static unsigned maxDepth() { return EnzymeMaxTypeDepth; } + + static int maxOffset() { return MaxTypeOffset; } + + template + static void depthLimitExceeded(const TreeT &Tree, Origin Orig, + const llvm::Twine &Detail) { + if (!EnzymeTypeWarning) + return; + if (CustomErrorHandler) { + CustomErrorHandler("TypeAnalysisDepthLimit", wrap(Orig), + ErrorType::TypeDepthExceeded, &Tree, nullptr, nullptr); + } else if (Orig) { + EmitWarning("TypeAnalysisDepthLimit", *Orig, *Orig, + " not handling more than ", EnzymeMaxTypeDepth, + " pointer lookups deep dt: ", Tree.str(), Detail.str()); + } else { + llvm::errs() << "not handling more than " << EnzymeMaxTypeDepth + << " pointer lookups deep dt:" << Tree.str() << Detail + << "\n"; + } + } +}; + +template class GenericTypeTree; + +using TypeTree = GenericTypeTree; typedef std::shared_ptr TypeResult; typedef std::map, ConcreteType> ConcreteTypeMapType; @@ -69,26 +100,38 @@ typedef std::map, const TypeResult> TypeTreeMapType; /// Class representing the underlying types of values as /// sequences of offsets to a ConcreteType -class TypeTree : public std::enable_shared_from_this { +/// +/// The class is templated over the handle of the floating point type, like +/// GenericConcreteType, so that it is shared between LLVM and MLIR +template +class GenericTypeTree + : public std::enable_shared_from_this> { + using Traits = TypeTreeTraits; + +public: + using ConcreteType = GenericConcreteType; + using ConcreteTypeMapType = std::map, ConcreteType>; + private: // mapping of known indices to type if one exists ConcreteTypeMapType mapping; std::vector minIndices; public: - TypeTree() {} - TypeTree(ConcreteType dat) { + GenericTypeTree() {} + GenericTypeTree(ConcreteType dat) { if (dat != ConcreteType(BaseType::Unknown)) { mapping.insert(std::pair, ConcreteType>({}, dat)); } } - static TypeTree parse(llvm::StringRef str, llvm::LLVMContext &ctx) { + template + static GenericTypeTree parse(llvm::StringRef str, ContextT &ctx) { using namespace llvm; assert(str[0] == '{'); str = str.substr(1); - TypeTree Result; + GenericTypeTree Result; while (true) { while (str[0] == ' ') str = str.substr(1); @@ -174,41 +217,37 @@ class TypeTree : public std::enable_shared_from_this { /// Lookup the underlying ConcreteType at a given offset sequence /// or Unknown if none exists - ConcreteType operator[](const std::vector Seq) const { - auto Found0 = mapping.find(Seq); - if (Found0 != mapping.end()) - return Found0->second; - size_t Len = Seq.size(); - if (Len == 0) + ConcreteType operator[](const std::vector &Seq) const { + auto Found = mapping.find(Seq); + if (Found != mapping.end()) + return Found->second; + if (Seq.empty()) return BaseType::Unknown; - std::vector> todo[2]; - todo[0].push_back({}); - int parity = 0; - for (size_t i = 0, Len = Seq.size(); i < Len - 1; ++i) { - for (auto prev : todo[parity]) { - prev.push_back(-1); - if (mapping.find(prev) != mapping.end()) - todo[1 - parity].push_back(prev); + std::vector> Prefixes = {{}}; + for (size_t i = 0; i + 1 < Seq.size(); ++i) { + std::vector> Next; + for (std::vector Cand : Prefixes) { + Cand.push_back(-1); + if (mapping.count(Cand)) + Next.push_back(Cand); if (Seq[i] != -1) { - prev.back() = Seq[i]; - if (mapping.find(prev) != mapping.end()) - todo[1 - parity].push_back(prev); + Cand.back() = Seq[i]; + if (mapping.count(Cand)) + Next.push_back(std::move(Cand)); } } - todo[parity].clear(); - parity = 1 - parity; + Prefixes = std::move(Next); } - size_t i = Len - 1; - for (auto prev : todo[parity]) { - prev.push_back(-1); - auto Found = mapping.find(prev); + for (std::vector Cand : Prefixes) { + Cand.push_back(-1); + auto Found = mapping.find(Cand); if (Found != mapping.end()) return Found->second; - if (Seq[i] != -1) { - prev.back() = Seq[i]; - Found = mapping.find(prev); + if (Seq.back() != -1) { + Cand.back() = Seq.back(); + Found = mapping.find(Cand); if (Found != mapping.end()) return Found->second; } @@ -236,18 +275,9 @@ class TypeTree : public std::enable_shared_from_this { bool insert(const std::vector Seq, ConcreteType CT, bool PointerIntSame = false) { size_t SeqSize = Seq.size(); - if (SeqSize > EnzymeMaxTypeDepth) { - if (EnzymeTypeWarning) { - if (CustomErrorHandler) { - CustomErrorHandler("TypeAnalysisDepthLimit", nullptr, - ErrorType::TypeDepthExceeded, this, nullptr, - nullptr); - } else - llvm::errs() << "not handling more than " << EnzymeMaxTypeDepth - << " pointer lookups deep dt:" << str() - << " adding v: " << to_string(Seq) << ": " << CT.str() - << "\n"; - } + if (SeqSize > Traits::maxDepth()) { + Traits::depthLimitExceeded( + *this, nullptr, " adding v: " + to_string(Seq) + ": " + CT.str()); return false; } if (SeqSize == 0) { @@ -373,7 +403,7 @@ class TypeTree : public std::enable_shared_from_this { (minIndices.size() <= SeqSize) ? minIndices.size() : SeqSize; for (size_t i = 0; i < minLen; i++) { if (minIndices[i] > Seq[i]) { - if (minIndices[i] > MaxTypeOffset) + if (minIndices[i] > Traits::maxOffset()) possibleDeletion = true; minIndices[i] = Seq[i]; } @@ -391,7 +421,7 @@ class TypeTree : public std::enable_shared_from_this { bool mustKeep = false; bool considerErase = false; for (int val : pair.first) { - if (val > MaxTypeOffset) { + if (val > Traits::maxOffset()) { if (val == minIndices[i]) { mustKeep = true; break; @@ -411,7 +441,7 @@ class TypeTree : public std::enable_shared_from_this { bool keep = false; bool considerErase = false; for (auto val : Seq) { - if (val > MaxTypeOffset) { + if (val > Traits::maxOffset()) { if (val == minIndices[i]) { keep = true; break; @@ -427,7 +457,9 @@ class TypeTree : public std::enable_shared_from_this { } /// How this TypeTree compares with another - bool operator<(const TypeTree &vd) const { return mapping < vd.mapping; } + bool operator<(const GenericTypeTree &vd) const { + return mapping < vd.mapping; + } /// Whether this TypeTree contains any information bool isKnown() const { @@ -458,8 +490,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Select only the Integer ConcreteTypes - TypeTree JustInt() const { - TypeTree vd; + GenericTypeTree JustInt() const { + GenericTypeTree vd; for (auto &pair : mapping) { if (pair.second == BaseType::Integer) { vd.insert(pair.first, pair.second); @@ -470,35 +502,21 @@ class TypeTree : public std::enable_shared_from_this { } /// Prepend an offset to all mappings - TypeTree Only(int Off, llvm::Instruction *orig) const { - TypeTree Result; + GenericTypeTree Only(int Off, typename Traits::Origin orig) const { + GenericTypeTree Result; Result.minIndices.reserve(1 + minIndices.size()); Result.minIndices.push_back(Off); for (auto midx : minIndices) Result.minIndices.push_back(midx); - if (Result.minIndices.size() > EnzymeMaxTypeDepth) { + if (Result.minIndices.size() > Traits::maxDepth()) { Result.minIndices.pop_back(); - if (EnzymeTypeWarning) { - if (CustomErrorHandler) { - CustomErrorHandler("TypeAnalysisDepthLimit", wrap(orig), - ErrorType::TypeDepthExceeded, this, nullptr, - nullptr); - } else if (orig) { - EmitWarning("TypeAnalysisDepthLimit", *orig, *orig, - " not handling more than ", EnzymeMaxTypeDepth, - " pointer lookups deep dt: ", str(), " only(", Off, ")"); - } else { - llvm::errs() << "not handling more than " << EnzymeMaxTypeDepth - << " pointer lookups deep dt:" << str() << " only(" - << Off << "): " - << "\n"; - } - } + Traits::depthLimitExceeded(*this, orig, + " only(" + llvm::Twine(Off) + ")"); } for (const auto &pair : mapping) { - if (pair.first.size() == EnzymeMaxTypeDepth) + if (pair.first.size() == Traits::maxDepth()) continue; std::vector Vec; Vec.reserve(pair.first.size() + 1); @@ -512,8 +530,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Peel off the outermost index at offset 0 - TypeTree Data0() const { - TypeTree Result; + GenericTypeTree Data0() const { + GenericTypeTree Result; for (const auto &pair : mapping) { if (pair.first.size() == 0) { @@ -555,8 +573,8 @@ class TypeTree : public std::enable_shared_from_this { /// Remove any mappings in the range [start, end) or [len, inf) /// This function has special handling for -1's - TypeTree Clear(size_t start, size_t end, size_t len) const { - TypeTree Result; + GenericTypeTree Clear(size_t start, size_t end, size_t len) const { + GenericTypeTree Result; // Note that below do insertion with the orIn operator // to force an error if there is an incompatible @@ -592,7 +610,7 @@ class TypeTree : public std::enable_shared_from_this { /// Select all submappings whose first index is in range [0, len) and remove /// the first index. This is the inverse of the `Only` operation - TypeTree Lookup(size_t len, const llvm::DataLayout &dl) const { + GenericTypeTree Lookup(size_t len, const llvm::DataLayout &dl) const { // Map of indices[1:] => ( End => possible Index[0] ) std::map, std::map>> staging; @@ -621,7 +639,7 @@ class TypeTree : public std::enable_shared_from_this { staging[next][pair.second].insert(pair.first[1]); } - TypeTree Result; + GenericTypeTree Result; for (auto &pair : staging) { auto &pnext = pair.first; for (auto &pair2 : pair.second) { @@ -784,8 +802,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Keep only pointers (or anything's) to a repeated value (represented by -1) - TypeTree KeepMinusOne(bool &legal) const { - TypeTree dat; + GenericTypeTree KeepMinusOne(bool &legal) const { + GenericTypeTree dat; for (const auto &pair : mapping) { @@ -826,7 +844,7 @@ class TypeTree : public std::enable_shared_from_this { size_t size = (dl.getTypeSizeInBits(val->getType()) + 7) / 8; - auto m0 = TypeTree::operator[]({0}); + auto m0 = GenericTypeTree::operator[]({0}); llvm::Type *flt = m0.isFloat(); if (!flt) { @@ -838,7 +856,7 @@ class TypeTree : public std::enable_shared_from_this { // walk bytewise until a concrete float type is found. size_t chunk = flt ? dl.getTypeSizeInBits(flt) / 8 : 1; for (size_t i = chunk; i < size; i += chunk) { - auto mx = TypeTree::operator[]({(int)i}); + auto mx = GenericTypeTree::operator[]({(int)i}); if (auto f2 = mx.isFloat()) { if (f2 != flt) { if (anythingIsFloat && !flt) { @@ -909,8 +927,8 @@ class TypeTree : public std::enable_shared_from_this { /// Replace mappings in the range in [offset, offset+maxSize] with those in // [addOffset, addOffset + maxSize]. In other words, select all mappings in // [offset, offset+maxSize] then add `addOffset` - TypeTree ShiftIndices(const llvm::DataLayout &dl, const int offset, - const int maxSize, size_t addOffset = 0) const { + GenericTypeTree ShiftIndices(const llvm::DataLayout &dl, const int offset, + const int maxSize, size_t addOffset = 0) const { // If we have no terms 1+ layer deep return the current result as a shift // won't change anything. This also makes the latercode simpler as it @@ -923,9 +941,9 @@ class TypeTree : public std::enable_shared_from_this { // will always result in an added variable (which would not be the case // on a size == 0). if (maxSize == 0) - return TypeTree(); + return GenericTypeTree(); - TypeTree Result; + GenericTypeTree Result; // The normal orIn / insert methods do collision checking, which is slow // (and presently O(n)). This is because an expansion of a -1 which could @@ -1044,7 +1062,7 @@ class TypeTree : public std::enable_shared_from_this { // If after moving this would not merit being kept for being a min index // or being within the max type offset, skip it. - if (next0 > MaxTypeOffset) { + if (next0 > Traits::maxOffset()) { bool minIndex = next0 == Result.minIndices[0]; if (!minIndex) for (size_t i = 1; i < pair.first.size(); i++) { @@ -1110,8 +1128,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Keep only mappings where the type is not an `Anything` - TypeTree PurgeAnything() const { - TypeTree Result; + GenericTypeTree PurgeAnything() const { + GenericTypeTree Result; Result.minIndices.reserve(minIndices.size()); for (const auto &pair : mapping) { if (pair.second == ConcreteType(BaseType::Anything)) @@ -1128,8 +1146,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Replace -1 with 0 - TypeTree ReplaceMinus() const { - TypeTree dat; + GenericTypeTree ReplaceMinus() const { + GenericTypeTree dat; for (const auto &pair : mapping) { if (pair.second == ConcreteType(BaseType::Anything)) continue; @@ -1152,8 +1170,8 @@ class TypeTree : public std::enable_shared_from_this { } /// Keep only mappings where the type is an `Anything` - TypeTree JustAnything() const { - TypeTree dat; + GenericTypeTree JustAnything() const { + GenericTypeTree dat; for (const auto &pair : mapping) { if (pair.second != ConcreteType(BaseType::Anything)) continue; @@ -1163,10 +1181,12 @@ class TypeTree : public std::enable_shared_from_this { } /// Chceck equality of two TypeTrees - bool operator==(const TypeTree &RHS) const { return mapping == RHS.mapping; } + bool operator==(const GenericTypeTree &RHS) const { + return mapping == RHS.mapping; + } /// Set this to another TypeTree, returning if this was changed - bool operator=(const TypeTree &RHS) { + bool operator=(const GenericTypeTree &RHS) { if (*this == RHS) return false; minIndices = RHS.minIndices; @@ -1312,7 +1332,8 @@ class TypeTree : public std::enable_shared_from_this { /// Set this to the logical or of itself and RHS, returning whether this value /// changed Setting `PointerIntSame` considers pointers and integers as /// equivalent If this is an illegal operation, `LegalOr` will be set to false - bool checkedOrIn(const TypeTree &RHS, bool PointerIntSame, bool &LegalOr) { + bool checkedOrIn(const GenericTypeTree &RHS, bool PointerIntSame, + bool &LegalOr) { // TODO detect recursive merge and simplify bool changed = false; @@ -1325,7 +1346,7 @@ class TypeTree : public std::enable_shared_from_this { /// Set this to the logical or of itself and RHS, returning whether this value /// changed Setting `PointerIntSame` considers pointers and integers as /// equivalent This function will error if doing an illegal Operation - bool orIn(const TypeTree &RHS, bool PointerIntSame) { + bool orIn(const GenericTypeTree &RHS, bool PointerIntSame) { bool Legal = true; bool Result = checkedOrIn(RHS, PointerIntSame, Legal); if (!Legal) { @@ -1356,14 +1377,14 @@ class TypeTree : public std::enable_shared_from_this { /// Set this to the logical or of itself and RHS, returning whether this value /// changed This assumes that pointers and integers are distinct This function /// will error if doing an illegal Operation - bool operator|=(const TypeTree &RHS) { + bool operator|=(const GenericTypeTree &RHS) { return orIn(RHS, /*PointerIntSame*/ false); } /// Set this to the logical and of itself and RHS, returning whether this /// value changed If this and RHS are incompatible at an index, the result /// will be BaseType::Unknown - bool andIn(const TypeTree &RHS) { + bool andIn(const GenericTypeTree &RHS) { bool changed = false; for (auto &pair : llvm::make_early_inc_range(mapping)) { @@ -1384,12 +1405,12 @@ class TypeTree : public std::enable_shared_from_this { /// Set this to the logical and of itself and RHS, returning whether this /// value changed If this and RHS are incompatible at an index, the result /// will be BaseType::Unknown - bool operator&=(const TypeTree &RHS) { return andIn(RHS); } + bool operator&=(const GenericTypeTree &RHS) { return andIn(RHS); } /// Set this to the logical `binop` of itself and RHS, using the Binop Op, /// returning true if this was changed. /// This function will error on an invalid type combination - bool binopIn(bool &Legal, const TypeTree &RHS, + bool binopIn(bool &Legal, const GenericTypeTree &RHS, llvm::BinaryOperator::BinaryOps Op) { bool changed = false; @@ -1474,7 +1495,7 @@ class TypeTree : public std::enable_shared_from_this { llvm::MDNode *toMD(llvm::LLVMContext &ctx) { llvm::SmallVector subMD; - std::map todo; + std::map todo; ConcreteType base(BaseType::Unknown); for (auto &pair : mapping) { if (pair.first.size() == 0) { @@ -1513,8 +1534,8 @@ class TypeTree : public std::enable_shared_from_this { } } - static TypeTree fromMD(llvm::MDNode *md) { - TypeTree ret; + static GenericTypeTree fromMD(llvm::MDNode *md) { + GenericTypeTree ret; std::vector off; ret.insertFromMD(md, off); return ret; diff --git a/enzyme/test/MLIR/ActivityAnalysis/nonfloat_result.mlir b/enzyme/test/MLIR/ActivityAnalysis/nonfloat_result.mlir new file mode 100644 index 00000000000..8a4e2fa0cd0 --- /dev/null +++ b/enzyme/test/MLIR/ActivityAnalysis/nonfloat_result.mlir @@ -0,0 +1,54 @@ +// RUN: %eopt --pass-pipeline="builtin.module(print-activity-analysis{dataflow=false annotate=true})" %s --split-input-file 2>&1 | FileCheck %s + +// An scf.for that only yields an integer but stores a float computation into +// memory must stay active: it writes memory, so the non-float result does not +// make it inactive. + +func.func @store_in_loop(%x: f64, %m: memref<10xf64>) -> index { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + %r = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %c0) -> (index) { + %y = arith.mulf %x, %x : f64 + memref.store %y, %m[%i] : memref<10xf64> + %n = arith.addi %acc, %c1 : index + scf.yield %n : index + } + return %r : index +} + +// CHECK: func.func @store_in_loop(%arg0: f64, %arg1: memref<10xf64>) -> index attributes {enzyme.arg_icv0 = false, enzyme.arg_icv1 = false, enzyme.ici = false} { +// CHECK: %0 = scf.for %arg2 = %c0 to %c10 step %c1 iter_args(%arg3 = %c0) -> (index) { +// CHECK-NEXT: %1 = arith.mulf %arg0, %arg0 {enzyme.ici = false, enzyme.res_icv0 = false} : f64 +// CHECK-NEXT: memref.store %1, %arg1[%arg2] {enzyme.ici = false} : memref<10xf64> +// CHECK-NEXT: %2 = arith.addi %arg3, %c1 {enzyme.ici = true, enzyme.res_icv0 = true} : index +// CHECK-NEXT: scf.yield {enzyme.ici = true} %2 : index +// CHECK-NEXT: } {enzyme.arg_icv0 = true, enzyme.arg_icv1 = true, enzyme.ici = false, enzyme.res_icv0 = true} + +// ----- + +// An scf.for that does not write memory and only yields an integer is +// inactive, even though its body computes on floats. + +func.func @readonly_loop(%x: f64) -> index { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + %r = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %c0) -> (index) { + %y = arith.mulf %x, %x : f64 + %t = arith.fptosi %y : f64 to i64 + %ti = arith.index_cast %t : i64 to index + %n = arith.addi %acc, %ti : index + scf.yield %n : index + } + return %r : index +} + +// CHECK: func.func @readonly_loop(%arg0: f64) -> index attributes {enzyme.arg_icv0 = false, enzyme.ici = true} { +// CHECK: %0 = scf.for %arg1 = %c0 to %c10 step %c1 iter_args(%arg2 = %c0) -> (index) { +// CHECK-NEXT: %1 = arith.mulf %arg0, %arg0 {enzyme.ici = true, enzyme.res_icv0 = true} : f64 +// CHECK-NEXT: %2 = arith.fptosi %1 {enzyme.ici = true, enzyme.res_icv0 = true} : f64 to i64 +// CHECK-NEXT: %3 = arith.index_cast %2 {enzyme.ici = true, enzyme.res_icv0 = true} : i64 to index +// CHECK-NEXT: %4 = arith.addi %arg2, %3 {enzyme.ici = true, enzyme.res_icv0 = true} : index +// CHECK-NEXT: scf.yield {enzyme.ici = true} %4 : index +// CHECK-NEXT: } {enzyme.arg_icv0 = true, enzyme.arg_icv1 = true, enzyme.ici = true, enzyme.res_icv0 = true}