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 //===----------------------------------------------------------------------===// 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,65 @@ 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. There are no restrictions placed on the shaped type or element + type of this attribute. + + 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">:$handle + ); + let builders = [ + // 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. + AttrBuilderWithInferredContext<(ins + "ShapedType":$type, "DenseResourceElementsHandle":$handle + )>, + // 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. + AttrBuilderWithInferredContext<(ins + "ShapedType":$type, "StringRef":$blobName, "AsmResourceBlob":$blob + )>, + // A builder that inserts a new resource using the provided blob. The handle + // of the inserted blob is used when building the attribute. + AttrBuilderWithInferredContext<(ins + "ShapedType":$type, "AsmResourceBlob":$blob + )>, + ]; +} + //===----------------------------------------------------------------------===// // 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.getHandle()); + 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,35 @@ 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(context, type, interface->insert(blobName, std::move(blob))); +} + +DenseResourceElementsAttr DenseResourceElementsAttr::get(ShapedType type, + AsmResourceBlob blob) { + // If we don't care about the resource name, just use a sensible default. + return get(type, "resource", std::move(blob)); +} + //===----------------------------------------------------------------------===// // 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/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; @@ -928,6 +933,38 @@ 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");