From 920daed78398efa8ed979de85564454e60d5cc3e Mon Sep 17 00:00:00 2001 From: River Riddle Date: Mon, 30 Jan 2023 22:42:14 -0800 Subject: [PATCH] [mlir][SubElements] Remove the ability to override implementations It's much cleaner and simpler to drive wacky configs via the AttrTypeSubElementHandler interface, instead of override. --- mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h | 6 +-- mlir/include/mlir/IR/AttrTypeSubElements.h | 30 +++++++++------ mlir/include/mlir/IR/StorageUniquerSupport.h | 33 ++-------------- mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp | 17 --------- mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h | 56 +++++++++++++++++++++++----- mlir/test/lib/Dialect/Test/TestTypes.h | 12 ------ 6 files changed, 70 insertions(+), 84 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h index a44448e..9ae0ba6 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h @@ -180,6 +180,7 @@ public: StringRef, bool); static LogicalResult verify(function_ref emitError, ArrayRef types, bool); + using Base::verify; /// Hooks for DataLayoutTypeInterface. Should not be called directly. Obtain a /// DataLayout instance and query it instead. @@ -197,11 +198,6 @@ public: LogicalResult verifyEntries(DataLayoutEntryListRef entries, Location loc) const; - - void walkImmediateSubElements(function_ref walkAttrsFn, - function_ref walkTypesFn) const; - Type replaceImmediateSubElements(ArrayRef replAttrs, - ArrayRef replTypes) const; }; //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/IR/AttrTypeSubElements.h b/mlir/include/mlir/IR/AttrTypeSubElements.h index 23b8d53..fe3f4cd 100644 --- a/mlir/include/mlir/IR/AttrTypeSubElements.h +++ b/mlir/include/mlir/IR/AttrTypeSubElements.h @@ -459,21 +459,29 @@ auto replaceImmediateSubElementsImpl(T derived, ArrayRef &replAttrs, // Otherwise, we need to replace any necessary sub-elements. } else { + // Functor used to build the replacement on success. + auto buildReplacement = [&](auto newKey, MLIRContext *ctx) { + if constexpr (is_tuple::value) { + return std::apply( + [&](auto &&...params) { + return constructSubElementReplacement( + ctx, std::forward(params)...); + }, + newKey); + } else { + return constructSubElementReplacement(ctx, newKey); + } + }; + AttrSubElementReplacements attrRepls(replAttrs); TypeSubElementReplacements typeRepls(replTypes); auto newKey = AttrTypeSubElementHandler::replace( key, attrRepls, typeRepls); - if constexpr (is_tuple::value) { - return std::apply( - [&](auto &&...params) { - return constructSubElementReplacement( - derived.getContext(), - std::forward(params)...); - }, - newKey); - } else { - return constructSubElementReplacement(derived.getContext(), newKey); - } + MLIRContext *ctx = derived.getContext(); + if constexpr (std::is_convertible_v) + return succeeded(newKey) ? buildReplacement(*newKey, ctx) : nullptr; + else + return buildReplacement(newKey, ctx); } } else { return derived; diff --git a/mlir/include/mlir/IR/StorageUniquerSupport.h b/mlir/include/mlir/IR/StorageUniquerSupport.h index 061bd67..333618b 100644 --- a/mlir/include/mlir/IR/StorageUniquerSupport.h +++ b/mlir/include/mlir/IR/StorageUniquerSupport.h @@ -127,38 +127,13 @@ public: }; } - /// Walk all of the immediately nested sub-attributes and sub-types. This - /// method does not recurse into sub elements. - void walkImmediateSubElements(function_ref walkAttrsFn, - function_ref walkTypesFn) const { - ::mlir::detail::walkImmediateSubElementsImpl( - *static_cast(this), walkAttrsFn, walkTypesFn); - } - - /// Replace the immediately nested sub-attributes and sub-types with those - /// provided. The order of the provided elements is derived from the order of - /// the elements returned by the callbacks of `walkImmediateSubElements`. The - /// element at index 0 would replace the very first attribute given by - /// `walkImmediateSubElements`. On success, the new instance with the values - /// replaced is returned. If replacement fails, nullptr is returned. - /// - /// Note that replacing the sub-elements of mutable types or attributes is - /// not currently supported by the interface. If an implementing type or - /// attribute is mutable, it should return `nullptr` if it has no mechanism - /// for replacing sub elements. - auto replaceImmediateSubElements(ArrayRef replAttrs, - ArrayRef replTypes) const { - return ::mlir::detail::replaceImmediateSubElementsImpl( - *static_cast(this), replAttrs, replTypes); - } - /// Returns a function that walks immediate sub elements of a given instance /// of the storage user. static auto getWalkImmediateSubElementsFn() { return [](auto instance, function_ref walkAttrsFn, function_ref walkTypesFn) { - cast(instance).walkImmediateSubElements(walkAttrsFn, - walkTypesFn); + ::mlir::detail::walkImmediateSubElementsImpl(cast(instance), + walkAttrsFn, walkTypesFn); }; } @@ -167,8 +142,8 @@ public: static auto getReplaceImmediateSubElementsFn() { return [](auto instance, ArrayRef replAttrs, ArrayRef replTypes) { - return cast(instance).replaceImmediateSubElements(replAttrs, - replTypes); + return ::mlir::detail::replaceImmediateSubElementsImpl( + cast(instance), replAttrs, replTypes); }; } diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp index ce28aa0..2e6bf8602 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp @@ -643,23 +643,6 @@ LogicalResult LLVMStructType::verifyEntries(DataLayoutEntryListRef entries, return mlir::success(); } -void LLVMStructType::walkImmediateSubElements( - function_ref walkAttrsFn, - function_ref walkTypesFn) const { - for (Type type : getBody()) - walkTypesFn(type); -} - -Type LLVMStructType::replaceImmediateSubElements( - ArrayRef replAttrs, ArrayRef replTypes) const { - if (isIdentified()) { - // TODO: It's not clear how we support replacing sub-elements of mutable - // types. - return nullptr; - } - return getLiteral(getContext(), replTypes, isPacked()); -} - //===----------------------------------------------------------------------===// // Vector types. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h b/mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h index d30e94c..2040d0a 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h +++ b/mlir/lib/Dialect/LLVMIR/IR/TypeDetail.h @@ -69,8 +69,9 @@ public: class Key { public: /// Constructs a key for an identified struct. - Key(StringRef name, bool opaque) - : name(name), identified(true), packed(false), opaque(opaque) {} + Key(StringRef name, bool opaque, ArrayRef types = std::nullopt) + : types(types), name(name), identified(true), packed(false), + opaque(opaque) {} /// Constructs a key for a literal struct. Key(ArrayRef types, bool packed) : types(types), identified(false), packed(packed), opaque(false) {} @@ -102,6 +103,13 @@ public: return types; } + /// Returns the list of type contained in an identified struct. + ArrayRef getIdentifiedStructBody() const { + assert(isIdentified() && + "requested struct body on a non-identified struct"); + return types; + } + /// Returns the hash value of the key. This combines various flags into a /// single value: the identified flag sets the first bit, and the packedness /// flag sets the second bit. Opacity bit is only used for construction and @@ -220,7 +228,7 @@ public: } /// Hook into the type uniquing infrastructure. - bool operator==(const KeyTy &other) const { return getKey() == other; }; + bool operator==(const KeyTy &other) const { return getAsKey() == other; }; static llvm::hash_code hashKey(const KeyTy &key) { return key.hashValue(); } static LLVMStructTypeStorage *construct(TypeStorageAllocator &allocator, const KeyTy &key) { @@ -251,6 +259,13 @@ public: return success(); } + /// Returns the key for the current storage. + Key getAsKey() const { + if (isIdentified()) + return Key(getIdentifier(), isOpaque(), getIdentifiedStructBody()); + return Key(getTypeList(), isPacked()); + } + private: /// Returns the number of elements in the key. unsigned keySize() const { @@ -271,13 +286,6 @@ private: llvm::Bitfield::set(identifiedBodySizeAndFlags, value); } - /// Returns the key for the current storage. - Key getKey() const { - if (isIdentified()) - return Key(getIdentifier(), isOpaque()); - return Key(getTypeList(), isPacked()); - } - /// Bitfield elements for `keyAndSizeFlags`: /// - bit 0: identified key flag; /// - bit 1: packed key flag; @@ -320,7 +328,35 @@ private: /// mutable flags. Must only be used through the Mutable* bitfields. unsigned identifiedBodySizeAndFlags = 0; }; +} // end namespace detail +} // end namespace LLVM + +/// Allow walking and replacing the subelements of a LLVMStructTypeStorage key. +template <> +struct AttrTypeSubElementHandler { + static void walk(const LLVM::detail::LLVMStructTypeStorage::Key ¶m, + AttrTypeImmediateSubElementWalker &walker) { + if (param.isIdentified()) + walker.walkRange(param.getIdentifiedStructBody()); + else + walker.walkRange(param.getTypeList()); + } + static FailureOr + replace(const LLVM::detail::LLVMStructTypeStorage::Key ¶m, + AttrSubElementReplacements &attrRepls, + TypeSubElementReplacements &typeRepls) { + // TODO: It's not clear how we support replacing sub-elements of mutable + // types. + if (param.isIdentified()) + return failure(); + + return LLVM::detail::LLVMStructTypeStorage::Key( + typeRepls.take_front(param.getTypeList().size()), param.isPacked()); + } +}; +namespace LLVM { +namespace detail { //===----------------------------------------------------------------------===// // LLVMTypeAndSizeStorage. //===----------------------------------------------------------------------===// diff --git a/mlir/test/lib/Dialect/Test/TestTypes.h b/mlir/test/lib/Dialect/Test/TestTypes.h index c6c8f58..c7d169d 100644 --- a/mlir/test/lib/Dialect/Test/TestTypes.h +++ b/mlir/test/lib/Dialect/Test/TestTypes.h @@ -146,18 +146,6 @@ public: /// Name/key getter. ::llvm::StringRef getName() { return getImpl()->name; } - - void walkImmediateSubElements( - ::llvm::function_ref walkAttrsFn, - ::llvm::function_ref walkTypesFn) const { - walkTypesFn(getBody()); - } - Type replaceImmediateSubElements(llvm::ArrayRef replAttrs, - llvm::ArrayRef replTypes) const { - // TODO: It's not clear how we support replacing sub-elements of mutable - // types. - return nullptr; - } }; } // namespace test -- 2.7.4