diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h --- a/mlir/include/mlir/IR/BuiltinAttributes.h +++ b/mlir/include/mlir/IR/BuiltinAttributes.h @@ -17,8 +17,12 @@ namespace mlir { class AffineMap; +class AsmResourceBlob; class BoolAttr; +class BuiltinDialect; class DenseIntElementsAttr; +template +struct DialectResourceBlobHandle; class FlatSymbolRefAttr; class FunctionType; class IntegerSet; @@ -731,6 +735,13 @@ return denseAttr && denseAttr.isSplat(); } }; + +//===----------------------------------------------------------------------===// +// DenseResourceElementsAttr +//===----------------------------------------------------------------------===// + +using DenseResourceElementsHandle = DialectResourceBlobHandle; + } // namespace mlir //===----------------------------------------------------------------------===// @@ -745,6 +756,9 @@ //===----------------------------------------------------------------------===// namespace mlir { +//===----------------------------------------------------------------------===// +// DenseArrayAttr + namespace detail { /// Base class for DenseArrayAttr that is instantiated and specialized for each /// supported element type below. @@ -794,6 +808,70 @@ using DenseF32ArrayAttr = detail::DenseArrayAttr; using DenseF64ArrayAttr = detail::DenseArrayAttr; +//===----------------------------------------------------------------------===// +// DenseResourceElementsAttr + +namespace detail { +/// Base class for DenseResourceElementsAttr that is instantiated and +/// specialized for each supported element type below. +template +class DenseResourceElementsAttrBase : public DenseResourceElementsAttr { +public: + using DenseResourceElementsAttr::DenseResourceElementsAttr; + + /// A builder that inserts a new resource using the provided blob. The handle + /// of the inserted blob is used when building the attribute. The provided + /// `blobName` is used as a hint for the key of the new handle for the `blob` + /// resource, but may be changed if necessary to ensure uniqueness during + /// insertion. + static DenseResourceElementsAttrBase get(ShapedType type, StringRef blobName, + AsmResourceBlob blob); + + /// Return the data of this attribute as an ArrayRef. + ArrayRef asArrayRef() const; + + /// Support for isa<>/cast<>. + static bool classof(Attribute attr); +}; + +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +extern template class DenseResourceElementsAttrBase; +} // namespace detail + +// Public names for all the supported DenseResourceElementsAttr. + +using DenseBoolResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseI8ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseI16ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseI32ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseI64ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseUI8ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseUI16ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseUI32ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseUI64ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseF32ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; +using DenseF64ResourceElementsAttr = + detail::DenseResourceElementsAttrBase; + //===----------------------------------------------------------------------===// // BoolAttr //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/IR/BuiltinAttributes.td b/mlir/include/mlir/IR/BuiltinAttributes.td --- a/mlir/include/mlir/IR/BuiltinAttributes.td +++ b/mlir/include/mlir/IR/BuiltinAttributes.td @@ -17,6 +17,7 @@ include "mlir/IR/AttrTypeBase.td" include "mlir/IR/BuiltinDialect.td" include "mlir/IR/BuiltinAttributeInterfaces.td" +include "mlir/IR/OpAsmInterface.td" include "mlir/IR/SubElementInterfaces.td" // TODO: Currently the attributes defined in this file are prefixed with @@ -420,6 +421,64 @@ let skipDefaultBuilders = 1; } +//===----------------------------------------------------------------------===// +// DenseResourceElementsAttr +//===----------------------------------------------------------------------===// + +def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements", [ + ElementsAttrInterface + ]> { + let summary = "An Attribute containing a dense multi-dimensional array " + "backed by a resource"; + let description = [{ + Syntax: + + ``` + dense-resource-elements-attribute ::= + `dense_resource` `<` resource-handle `>` `:` shaped-type + ``` + + A dense resource elements attribute is an elements attribute backed by a + handle to a builtin dialect resource containing a densely packed array of + values. This class provides the low-level attribute, which should only be + interacted with in very generic terms, actual access to the underlying + resource data is intended to be managed through one of the subclasses, such + as; `DenseBoolResourceElementsAttr`, `DenseUI64ResourceElementsAttr`, + `DenseI32ArrayAttr`, `DenseF32ArrayAttr`, `DenseF64ArrayAttr`, etc. + + Examples: + + ```mlir + // A tensor referencing a builtin dialect resource, `resource_1`, with two + // unsigned i32 elements. + dense_resource : tensor<2xui32> + ``` + }]; + let parameters = (ins + AttributeSelfTypeParameter<"", "ShapedType">:$type, + ResourceHandleParameter<"DenseResourceElementsHandle">:$rawHandle + ); + let builders = [ + AttrBuilderWithInferredContext<(ins + "ShapedType":$type, "DenseResourceElementsHandle":$handle + )> + ]; + let extraClassDeclaration = [{ + protected: + /// A builder that inserts a new resource using the provided blob. The handle + /// of the inserted blob is used when building the attribute. The provided + /// `blobName` is used as a hint for the key of the new handle for the + /// `blob` resource, but may be changed if necessary to ensure uniqueness + /// during insertion. + static DenseResourceElementsAttr get( + ShapedType type, StringRef blobName, AsmResourceBlob blob + ); + + public: + }]; + let skipDefaultBuilders = 1; +} + //===----------------------------------------------------------------------===// // DictionaryAttr //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -20,6 +20,7 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/DialectResourceBlobManager.h" #include "mlir/IR/IntegerSet.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/OpImplementation.h" @@ -1896,6 +1897,10 @@ os << " "; denseArrayAttr.printWithoutBraces(os); os << "]"; + } else if (auto resourceAttr = attr.dyn_cast()) { + os << "dense_resource<"; + printResourceHandle(resourceAttr.getRawHandle()); + os << ">"; } else if (auto locAttr = attr.dyn_cast()) { printLocation(locAttr); } else { diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp --- a/mlir/lib/IR/BuiltinAttributes.cpp +++ b/mlir/lib/IR/BuiltinAttributes.cpp @@ -11,6 +11,7 @@ #include "mlir/IR/AffineMap.h" #include "mlir/IR/BuiltinDialect.h" #include "mlir/IR/Dialect.h" +#include "mlir/IR/DialectResourceBlobManager.h" #include "mlir/IR/IntegerSet.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/Operation.h" @@ -1552,6 +1553,130 @@ attr.getType().cast().getElementType().isIntOrIndex(); } +//===----------------------------------------------------------------------===// +// DenseResourceElementsAttr +//===----------------------------------------------------------------------===// + +DenseResourceElementsAttr +DenseResourceElementsAttr::get(ShapedType type, + DenseResourceElementsHandle handle) { + return Base::get(type.getContext(), type, handle); +} + +DenseResourceElementsAttr DenseResourceElementsAttr::get(ShapedType type, + StringRef blobName, + AsmResourceBlob blob) { + // Extract the builtin dialect resource manager from context. + MLIRContext *context = type.getContext(); + BuiltinDialect *dialect = context->getLoadedDialect(); + auto *interface = dialect->getRegisteredInterface< + ResourceBlobManagerDialectInterfaceBase>(); + + // Construct a handle by inserting a new resource using the provided blob. + return get(type, interface->insert(blobName, std::move(blob))); +} + +//===----------------------------------------------------------------------===// +// DenseResourceElementsAttrBase + +namespace { +/// Instantiations of this class provide utilities for interacting with native +/// data types in the context of DenseResourceElementsAttr. +template +struct DenseResourceAttrUtil; +template +struct DenseResourceElementsAttrIntUtil { + static bool checkElementType(Type eltType) { + IntegerType type = eltType.dyn_cast(); + if (!type || type.getWidth() != width) + return false; + return isSigned ? !type.isUnsigned() : !type.isSigned(); + } +}; +template <> +struct DenseResourceAttrUtil { + static bool checkElementType(Type eltType) { + return eltType.isSignlessInteger(1); + } +}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<8, true> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<8, false> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<16, true> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<16, false> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<32, true> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<32, false> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<64, true> {}; +template <> +struct DenseResourceAttrUtil + : public DenseResourceElementsAttrIntUtil<64, false> {}; +template <> +struct DenseResourceAttrUtil { + static bool checkElementType(Type eltType) { return eltType.isF32(); } +}; +template <> +struct DenseResourceAttrUtil { + static bool checkElementType(Type eltType) { return eltType.isF64(); } +}; +} // namespace + +template +DenseResourceElementsAttrBase +DenseResourceElementsAttrBase::get(ShapedType type, StringRef blobName, + AsmResourceBlob blob) { + // Check that the blob is in the form we were expecting. + assert(blob.getDataAlignment() == alignof(T) && + "alignment mismatch between expected alignment and blob alignment"); + assert(((blob.getData().size() % sizeof(T)) == 0) && + "size mismatch between expected element width and blob size"); + assert(DenseResourceAttrUtil::checkElementType(type.getElementType()) && + "invalid shape element type for provided type `T`"); + return DenseResourceElementsAttr::get(type, blobName, std::move(blob)) + .template cast>(); +} + +template +ArrayRef DenseResourceElementsAttrBase::asArrayRef() const { + return this->getRawHandle().getBlob().template getDataAs(); +} + +template +bool DenseResourceElementsAttrBase::classof(Attribute attr) { + auto resourceAttr = attr.dyn_cast(); + return resourceAttr && DenseResourceAttrUtil::checkElementType( + resourceAttr.getElementType()); +} + +namespace mlir { +namespace detail { +// Explicit instantiation for all the supported DenseResourceElementsAttr. +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +template class DenseResourceElementsAttrBase; +} // namespace detail +} // namespace mlir + //===----------------------------------------------------------------------===// // OpaqueElementsAttr //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp --- a/mlir/lib/IR/BuiltinDialect.cpp +++ b/mlir/lib/IR/BuiltinDialect.cpp @@ -16,20 +16,34 @@ #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/DialectResourceBlobManager.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/PatternMatch.h" using namespace mlir; //===----------------------------------------------------------------------===// -// Builtin Dialect +// TableGen'erated dialect //===----------------------------------------------------------------------===// #include "mlir/IR/BuiltinDialect.cpp.inc" +//===----------------------------------------------------------------------===// +// BuiltinBlobManagerInterface +//===----------------------------------------------------------------------===// + +using BuiltinBlobManagerInterface = + ResourceBlobManagerDialectInterfaceBase; + +//===----------------------------------------------------------------------===// +// BuiltinOpAsmDialectInterface +//===----------------------------------------------------------------------===// + namespace { struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface { - using OpAsmDialectInterface::OpAsmDialectInterface; + BuiltinOpAsmDialectInterface(Dialect *dialect, + BuiltinBlobManagerInterface &mgr) + : OpAsmDialectInterface(dialect), blobManager(mgr) {} AliasResult getAlias(Attribute attr, raw_ostream &os) const override { if (attr.isa()) { @@ -56,6 +70,38 @@ } return AliasResult::NoAlias; } + + //===------------------------------------------------------------------===// + // Resources + //===------------------------------------------------------------------===// + + std::string + getResourceKey(const AsmDialectResourceHandle &handle) const override { + return cast(handle).getKey().str(); + } + FailureOr + declareResource(StringRef key) const final { + return blobManager.insert(key); + } + LogicalResult parseResource(AsmParsedResourceEntry &entry) const final { + FailureOr blob = entry.parseAsBlob(); + if (failed(blob)) + return failure(); + + // Update the blob for this entry. + blobManager.update(entry.getKey(), std::move(*blob)); + return success(); + } + void + buildResources(Operation *op, + const SetVector &referencedResources, + AsmResourceBuilder &provider) const final { + blobManager.buildResources(provider, referencedResources.getArrayRef()); + } + +private: + /// The blob manager for the dialect. + BuiltinBlobManagerInterface &blobManager; }; } // namespace @@ -67,7 +113,9 @@ #define GET_OP_LIST #include "mlir/IR/BuiltinOps.cpp.inc" >(); - addInterfaces(); + + auto &blobInterface = addInterface(); + addInterface(blobInterface); } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/DialectResourceBlobManager.cpp b/mlir/lib/IR/DialectResourceBlobManager.cpp --- a/mlir/lib/IR/DialectResourceBlobManager.cpp +++ b/mlir/lib/IR/DialectResourceBlobManager.cpp @@ -56,7 +56,7 @@ Twine(nameCounter++).toVector(nameStorage); // Try inserting with the new name. - if (BlobEntry *entry = tryInsertion(name)) + if (BlobEntry *entry = tryInsertion(nameStorage)) return *entry; nameStorage.resize(name.size() + 1); } while (true); diff --git a/mlir/lib/Parser/AttributeParser.cpp b/mlir/lib/Parser/AttributeParser.cpp --- a/mlir/lib/Parser/AttributeParser.cpp +++ b/mlir/lib/Parser/AttributeParser.cpp @@ -14,9 +14,10 @@ #include "AsmParserImpl.h" #include "mlir/IR/AffineMap.h" +#include "mlir/IR/BuiltinDialect.h" #include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Dialect.h" #include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/DialectResourceBlobManager.h" #include "mlir/IR/IntegerSet.h" #include "mlir/Parser/AsmParserState.h" #include "llvm/ADT/StringExtras.h" @@ -97,6 +98,10 @@ case Token::kw_dense: return parseDenseElementsAttr(type); + // Parse a dense resource elements attribute. + case Token::kw_dense_resource: + return parseDenseResourceElementsAttr(type); + // Parse a dictionary attribute. case Token::l_brace: { NamedAttrList elements; @@ -241,6 +246,7 @@ case Token::kw_affine_map: case Token::kw_affine_set: case Token::kw_dense: + case Token::kw_dense_resource: case Token::kw_false: case Token::kw_loc: case Token::kw_opaque: @@ -928,6 +934,39 @@ return literalParser.getAttr(loc, type); } +Attribute Parser::parseDenseResourceElementsAttr(Type attrType) { + auto loc = getToken().getLoc(); + consumeToken(Token::kw_dense_resource); + if (parseToken(Token::less, "expected '<' after 'dense_resource'")) + return nullptr; + + // Parse the resource handle. + FailureOr rawHandle = + parseResourceHandle(getContext()->getLoadedDialect()); + if (failed(rawHandle) || parseToken(Token::greater, "expected '>'")) + return nullptr; + + auto *handle = dyn_cast(&*rawHandle); + if (!handle) + return emitError(loc, "invalid `dense_resource` handle type"), nullptr; + + // Parse the type of the attribute if the user didn't provide one. + SMLoc typeLoc = loc; + if (!attrType) { + typeLoc = getToken().getLoc(); + if (parseToken(Token::colon, "expected ':'") || !(attrType = parseType())) + return nullptr; + } + + ShapedType shapedType = attrType.dyn_cast(); + if (!shapedType) { + emitError(typeLoc, "`dense_resource` expected a shaped type"); + return nullptr; + } + + return DenseResourceElementsAttr::get(shapedType, *handle); +} + /// Parse an opaque elements attribute. Attribute Parser::parseOpaqueElementsAttr(Type attrType) { SMLoc loc = getToken().getLoc(); diff --git a/mlir/lib/Parser/Parser.h b/mlir/lib/Parser/Parser.h --- a/mlir/lib/Parser/Parser.h +++ b/mlir/lib/Parser/Parser.h @@ -160,6 +160,7 @@ /// Parse a handle to a dialect resource within the assembly format. FailureOr parseResourceHandle(const OpAsmDialectInterface *dialect, StringRef &name); + FailureOr parseResourceHandle(Dialect *dialect); //===--------------------------------------------------------------------===// // Type Parsing @@ -272,6 +273,9 @@ Attribute parseDenseElementsAttr(Type attrType); ShapedType parseElementsLiteralType(Type type); + /// Parse a dense resource elements attribute. + Attribute parseDenseResourceElementsAttr(Type attrType); + /// Parse a DenseArrayAttr. Attribute parseDenseArrayAttr(); diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp --- a/mlir/lib/Parser/Parser.cpp +++ b/mlir/lib/Parser/Parser.cpp @@ -342,6 +342,17 @@ return entry.second; } +FailureOr +Parser::parseResourceHandle(Dialect *dialect) { + const auto *interface = dyn_cast_or_null(dialect); + if (!interface) { + return emitError() << "dialect '" << dialect->getNamespace() + << "' does not expect resource handles"; + } + StringRef resourceName; + return parseResourceHandle(interface, resourceName); +} + //===----------------------------------------------------------------------===// // Code Completion diff --git a/mlir/lib/Parser/TokenKinds.def b/mlir/lib/Parser/TokenKinds.def --- a/mlir/lib/Parser/TokenKinds.def +++ b/mlir/lib/Parser/TokenKinds.def @@ -87,6 +87,7 @@ TOK_KEYWORD(ceildiv) TOK_KEYWORD(complex) TOK_KEYWORD(dense) +TOK_KEYWORD(dense_resource) TOK_KEYWORD(f16) TOK_KEYWORD(f32) TOK_KEYWORD(f64) diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp --- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp +++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp @@ -704,8 +704,9 @@ /// Signal a completion for an attribute. void completeAttribute(const llvm::StringMap &aliases) override { - appendSimpleCompletions({"affine_set", "affine_map", "dense", "false", - "loc", "opaque", "sparse", "true", "unit"}, + appendSimpleCompletions({"affine_set", "affine_map", "dense", + "dense_resource", "false", "loc", "opaque", + "sparse", "true", "unit"}, lsp::CompletionItemKind::Field, /*sortText=*/"1"); diff --git a/mlir/test/IR/invalid-builtin-attributes.mlir b/mlir/test/IR/invalid-builtin-attributes.mlir --- a/mlir/test/IR/invalid-builtin-attributes.mlir +++ b/mlir/test/IR/invalid-builtin-attributes.mlir @@ -519,3 +519,23 @@ "J// ----- " // expected-error {{expected}} + +// ----- + +// expected-error@+1 {{expected '<' after 'dense_resource'}} +#attr = dense_resource> + +// ----- + +// expected-error@+1 {{expected '>'}} +#attr = dense_resource + +// ----- + +// expected-error@+1 {{`dense_resource` expected a shaped type}} +#attr = dense_resource : i32 diff --git a/mlir/test/IR/invalid-file-metadata.mlir b/mlir/test/IR/invalid-file-metadata.mlir --- a/mlir/test/IR/invalid-file-metadata.mlir +++ b/mlir/test/IR/invalid-file-metadata.mlir @@ -59,10 +59,10 @@ // ----- -// expected-error@+4 {{unknown 'resource' key 'unknown_entry' for dialect 'builtin'}} +// expected-error@+4 {{unknown 'resource' key 'unknown_entry' for dialect 'ml_program'}} {-# dialect_resources: { - builtin: { + ml_program: { unknown_entry: "foo" } }