[mlir] Add loaded URI attribute type #67276
Conversation
Add attribute type that behaves like a regular ElementsAttr with the exception that it is printed and parsed by using a URI (and offset and size within) and references within URI. This allows for an external to context and IR constant. It is still a constant, so its not for parameter case nor an extern constant. The offset and size allows for having a single file with multiple values inside (and for some custom formats that are mmap'able this ends up being sufficient if the layout matches). While the nested references allows for usage in formats that have an indexing table (e.g., MLIR bytecode). These differ as to when considered: the offset & size is used to dictate which part to load into a memory buffer, while the references are for once the file is loaded. In effect this is an attribute whose value doesn't get printed and that is usable where the external file will be available. These are also usable in cases where ElementsAttr is supported, such as ml_program globals or the like where it is not immediately a constant --- such usage would allow for some final level of customization/downstream adaption to decide what to inline when (ideally the correct ops would be created from the start, either with URI, with loaded URI or constant with loaded URI, so that frontend dictates these). The other potential use is in the elide large attribute case, where it could be outlined to files now and you get the small IR while having exactly the same behaviour when passed through again (while being friendly on editor/human reader). Note: this is a minimal initial version, itdoes not yet add support for different URI nor file types, nor lazy loading.
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir ChangesAdd attribute type that behaves like a regular ElementsAttr with the exception that it is printed and parsed by using a URI (and offset and size within) and references within URI. This allows for an external to context and IR constant. It is still a constant, so its not for parameter case nor an extern constant. The offset and size allows for having a single file with multiple values inside (and for some custom formats that are mmap'able this ends up being sufficient if the layout matches). While the nested references allows for usage in formats that have an indexing table (e.g., MLIR bytecode). These differ as to when considered: the offset & size is used to dictate which part to load into a memory buffer, while the references are for once the file is loaded. In effect this is an attribute whose value doesn't get printed and that is usable where the external file will be available. These are also usable in cases where ElementsAttr is supported, such as ml_program globals or the like where it is not immediately a constant --- such usage would allow for some final level of customization/downstream adaption to decide what to inline when (ideally the correct ops would be created from the start, either with URI, with loaded URI or constant with loaded URI, so that frontend dictates these). The other potential use is in the elide large attribute case, where it could be outlined to files now and you get the small IR while having exactly the same behaviour when passed through again (while being friendly on editor/human reader). Note: this is a minimal initial version, itdoes not yet add support for different URI nor file types, nor lazy loading. Patch is 20.52 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/67276.diff 9 Files Affected:
diff --git a/mlir/docs/Dialects/Builtin.md b/mlir/docs/Dialects/Builtin.md
index 0a9b7ae8919b55e..c6c9f3206175ca6 100644
--- a/mlir/docs/Dialects/Builtin.md
+++ b/mlir/docs/Dialects/Builtin.md
@@ -15,7 +15,7 @@ extensible by design, any potential additions are heavily scrutinized.
[include "Dialects/BuiltinAttributes.md"]
-## Location Attributes
+### Location Attributes
A subset of the builtin attribute values correspond to
[source locations](../Diagnostics.md/#source-locations), that may be attached to
@@ -23,7 +23,7 @@ Operations.
[include "Dialects/BuiltinLocationAttributes.md"]
-## DistinctAttribute
+### DistinctAttribute
A DistinctAttribute associates an attribute with a unique identifier.
As a result, multiple DistinctAttribute instances may point to the same
@@ -54,6 +54,23 @@ identifier, which can be used to mark groups of operations that share a
common property. For example, groups of aliasing memory operations may be
marked using one DistinctAttribute instance per alias group.
+### LoadedURIDenseResourceAttr
+
+A LoadedURIDenseResourceAttr is an attribute that implements the ElementsAttr
+interface, prints & parses as URI, optional offset, size, alignment and nested
+symbol reference paramters while keeping the value represented in memory (but
+outside the context). This allows for having constants whose value are outside
+the context and (during roundtrip) the IR, but without any semantic differences.
+
+```
+luri ::= `luri` `<` string-literal
+ (`,` `::` symbol-ref-attribute)?
+ (`,` `byte_offset` `=` integer-literal)?
+ (`,` `byte_size` `=` integer-literal)?
+ (`,` `alignment` `=` integer-literal)?
+ `>`
+```
+
## Operations
[include "Dialects/BuiltinOps.md"]
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index c8161604aad3503..508a3e6e7da7317 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -1046,6 +1046,92 @@ class DistinctAttr
static DistinctAttr create(Attribute referencedAttr);
};
+//===----------------------------------------------------------------------===//
+// LoadedURIDenseResourceAttr
+//===----------------------------------------------------------------------===//
+
+namespace detail {
+struct LoadedURIDenseResourceAttrStorage;
+} // namespace detail
+
+/// An attribute that parses/prints as a identifier (file & symbol) while in
+/// memory representing the loaded attribute.
+
+class LoadedURIDenseResourceAttr
+ : public Attribute::AttrBase<LoadedURIDenseResourceAttr, Attribute,
+ detail::LoadedURIDenseResourceAttrStorage,
+ TypedAttr::Trait, ElementsAttr::Trait> {
+public:
+ using Base::Base;
+ using Base::getChecked;
+
+ static LoadedURIDenseResourceAttr
+ get(StringRef uri, Type type, std::optional<int64_t> alignment,
+ std::optional<int64_t> offset, std::optional<int64_t> size,
+ ArrayRef<FlatSymbolRefAttr> nestedReferences = {});
+
+ static LoadedURIDenseResourceAttr
+ getChecked(function_ref<::mlir::InFlightDiagnostic()> emitError,
+ StringRef uri, Type type, std::optional<int64_t> alignment,
+ std::optional<int64_t> offset, std::optional<int64_t> size,
+ ArrayRef<FlatSymbolRefAttr> nestedReferences = {});
+
+ /// The set of data types that can be iterated by this attribute.
+ // TODO: Expand as needed; consider maybe reuse with DenseElementsAttr.
+ using ContiguousIterableTypesT = std::tuple<>;
+ using NonContiguousIterableTypesT = std::tuple<
+ // Float types.
+ APFloat>;
+
+ /// Returns the loaded value.
+ ArrayRef<char> getRawData() const;
+
+ /// Returns the underlying data as an array of the given type. This is an
+ /// inherently unsafe operation, and should only be used when the data is
+ /// known to be of the correct type.
+ template <typename T>
+ ArrayRef<T> getDataAs() const {
+ auto data = getRawData();
+ return llvm::ArrayRef<T>((const T *)data.data(), data.size() / sizeof(T));
+ }
+ /// Provide a `try_value_begin_impl` to enable iteration within
+ /// ElementsAttr.
+ auto try_value_begin_impl(OverloadToken<llvm::APFloat>) const {
+ auto it = llvm::map_range(getDataAs<float>(), [=](float value) {
+ return llvm::APFloat(value);
+ }).begin();
+ return FailureOr<decltype(it)>(std::move(it));
+ }
+
+ // TODO: Add member functions to register URI/file type handlers.
+
+ /// Return the URI loaded.
+ StringRef getURI() const;
+
+ /// Return the Type of the data.
+ ShapedType getType() const;
+
+ /// Return the nested references inside the URI loaded.
+ ArrayRef<FlatSymbolRefAttr> getNestedReferences() const;
+
+ /// Return the alignment the resource is stored with.
+ // TODO: Should this be exposed?
+ std::optional<int64_t> getAlignment() const;
+
+ /// Return the byte offset at which entry was in URI.
+ std::optional<int64_t> getByteOffset() const;
+
+ /// Return the byte size at which entry was in URI.
+ std::optional<int64_t> getByteSize() const;
+
+ static LogicalResult
+ verify(function_ref<::mlir::InFlightDiagnostic()> emitError, StringRef uri,
+ Type type, std::optional<int64_t> alignment,
+ std::optional<int64_t> offset, std::optional<int64_t> size,
+ ArrayRef<FlatSymbolRefAttr> nestedReferences,
+ std::optional<FailureOr<DenseResourceElementsHandle>> handle);
+};
+
//===----------------------------------------------------------------------===//
// StringAttr
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 3437ac9addc5ff6..6cde5e2574a795f 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -148,6 +148,10 @@ Attribute Parser::parseAttribute(Type type) {
return locAttr;
}
+ // Parse a loaded URI elements attribute.
+ case Token::kw_luri:
+ return parseLoadedURIElementsAttr(type);
+
// Parse a sparse elements attribute.
case Token::kw_sparse:
return parseSparseElementsAttr(type);
@@ -1065,6 +1069,110 @@ ShapedType Parser::parseElementsLiteralType(Type type) {
return sType;
}
+/// Parse a sparse elements attribute.
+Attribute Parser::parseLoadedURIElementsAttr(Type attrType) {
+ SMLoc loc = getToken().getLoc();
+ consumeToken(Token::kw_luri);
+ if (parseToken(Token::less, "expected '<' after 'luri'"))
+ return nullptr;
+
+ auto fileName = getToken().getStringValue();
+ if (parseToken(Token::string, "expected URI"))
+ return nullptr;
+
+ // Parse any nested references.
+ std::vector<FlatSymbolRefAttr> nestedRefs;
+ while (consumeIf(Token::colon)) {
+ // Check for the '::' prefix.
+ if (parseToken(Token::colon, "ex[ected '::' for nested reference"))
+ return nullptr;
+ // Parse the reference itself.
+ auto curLoc = getToken().getLoc();
+ if (getToken().isNot(Token::at_identifier)) {
+ emitError(curLoc, "expected nested symbol reference identifier");
+ return nullptr;
+ }
+
+ std::string nameStr = getToken().getSymbolReference();
+ consumeToken(Token::at_identifier);
+ nestedRefs.push_back(SymbolRefAttr::get(getContext(), nameStr));
+ }
+
+ std::optional<int64_t> alignment;
+ std::optional<int64_t> byteOffset;
+ std::optional<int64_t> byteSize;
+
+ StringRef keyword;
+ for (auto curLoc = getToken().getLoc();
+ consumeIf(Token::comma) && !parseOptionalKeyword(&keyword);) {
+ curLoc = getToken().getLoc();
+
+ if (keyword == "alignment") {
+ if (alignment.has_value()) {
+ emitError("multiple alignment values provided");
+ return nullptr;
+ }
+ if (parseToken(Token::equal, "expected '=' after alignment"))
+ return nullptr;
+ APInt val;
+ if (auto res = parseOptionalInteger(val);
+ !res.has_value() || failed(*res)) {
+ emitError("missing or invalid alignment");
+ return nullptr;
+ }
+ alignment = val.getSExtValue();
+ continue;
+ }
+
+ if (keyword == "byte_offset") {
+ if (byteOffset.has_value()) {
+ emitError("multiple offset values provided");
+ return nullptr;
+ }
+ if (parseToken(Token::equal, "expected '=' after offset"))
+ return nullptr;
+ APInt val;
+ if (auto res = parseOptionalInteger(val);
+ !res.has_value() || failed(*res)) {
+ emitError("missing or invalid offset");
+ return nullptr;
+ }
+ byteOffset = val.getSExtValue();
+ continue;
+ }
+
+ if (keyword == "byte_size") {
+ if (byteSize.has_value()) {
+ emitError("multiple size values provided");
+ return nullptr;
+ }
+ if (parseToken(Token::equal, "expected '=' after size"))
+ return nullptr;
+ APInt val;
+ if (auto res = parseOptionalInteger(val);
+ !res.has_value() || failed(*res)) {
+ emitError("missing or invalid size");
+ return nullptr;
+ }
+ byteSize = val.getSExtValue();
+ continue;
+ }
+
+ emitError(curLoc, "unexpected keyword: ") << keyword;
+ }
+
+ if (parseToken(Token::greater, "expected '>'"))
+ return nullptr;
+
+ auto type = parseElementsLiteralType(attrType);
+ if (!type)
+ return nullptr;
+
+ // Build the sparse elements attribute by the indices and values.
+ return getChecked<LoadedURIDenseResourceAttr>(
+ loc, fileName, type, alignment, byteOffset, byteSize, nestedRefs);
+}
+
/// Parse a sparse elements attribute.
Attribute Parser::parseSparseElementsAttr(Type attrType) {
SMLoc loc = getToken().getLoc();
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 01c55f97a08c2ce..dce9b4a758cf5b7 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -273,6 +273,9 @@ class Parser {
/// Parse a DenseArrayAttr.
Attribute parseDenseArrayAttr(Type type);
+ /// Parse a loaded URI elements attribute.
+ Attribute parseLoadedURIElementsAttr(Type type);
+
/// Parse a sparse elements attribute.
Attribute parseSparseElementsAttr(Type attrType);
diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def
index 297e07459453013..7c867eeaddb7b38 100644
--- a/mlir/lib/AsmParser/TokenKinds.def
+++ b/mlir/lib/AsmParser/TokenKinds.def
@@ -106,6 +106,7 @@ TOK_KEYWORD(for)
TOK_KEYWORD(func)
TOK_KEYWORD(index)
TOK_KEYWORD(loc)
+TOK_KEYWORD(luri)
TOK_KEYWORD(max)
TOK_KEYWORD(memref)
TOK_KEYWORD(min)
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 7b0da30541b16a4..1872a85cc00e280 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -2311,6 +2311,20 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
os << ">";
} else if (auto locAttr = llvm::dyn_cast<LocationAttr>(attr)) {
printLocation(locAttr);
+ } else if (auto luriAttr = llvm::dyn_cast<LoadedURIDenseResourceAttr>(attr)) {
+ os << "luri<";
+ printEscapedString(luriAttr.getURI());
+ for (FlatSymbolRefAttr nestedRef : luriAttr.getNestedReferences()) {
+ os << "::";
+ printSymbolReference(nestedRef.getValue(), os);
+ }
+ if (auto alignment = luriAttr.getAlignment())
+ os << ", alignment = " << alignment;
+ if (auto offset = luriAttr.getByteOffset())
+ os << ", byte_offset = " << offset;
+ if (auto size = luriAttr.getByteSize())
+ os << ", byte_size = " << size;
+ os << ">";
} else {
llvm::report_fatal_error("Unknown builtin attribute");
}
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 5f1129326f4f772..2a043b668806c98 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -21,6 +21,7 @@
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Endian.h"
+#include "llvm/Support/SourceMgr.h"
#include <optional>
using namespace mlir;
@@ -42,7 +43,7 @@ void BuiltinDialect::registerAttributes() {
#define GET_ATTRDEF_LIST
#include "mlir/IR/BuiltinAttributes.cpp.inc"
>();
- addAttributes<DistinctAttr>();
+ addAttributes<DistinctAttr, LoadedURIDenseResourceAttr>();
}
//===----------------------------------------------------------------------===//
@@ -1761,6 +1762,193 @@ Attribute DistinctAttr::getReferencedAttr() const {
return getImpl()->referencedAttr;
}
+//===----------------------------------------------------------------------===//
+// LoadedURIDenseResourceAttr
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace detail {
+struct LoadedURIDenseResourceAttrStorage : public ::mlir::AttributeStorage {
+ using KeyTy =
+ std::tuple<StringRef, Type, std::optional<int64_t>,
+ std::optional<int64_t>, std::optional<int64_t>,
+ ArrayRef<FlatSymbolRefAttr>,
+ std::optional<FailureOr<DenseResourceElementsHandle>>>;
+ LoadedURIDenseResourceAttrStorage(
+ StringRef uri, Type type, std::optional<int64_t> alignment,
+ std::optional<int64_t> offset, std::optional<int64_t> size,
+ ArrayRef<FlatSymbolRefAttr> nestedReferences,
+ std::optional<FailureOr<DenseResourceElementsHandle>> rawHandle)
+ : uri(uri), type(type), nestedReferences(nestedReferences),
+ alignment(alignment), offset(offset), size(size), rawHandle(rawHandle) {
+ }
+
+ KeyTy getAsKey() const {
+ return KeyTy(uri, type, alignment, offset, size, nestedReferences,
+ rawHandle);
+ }
+
+ bool operator==(const KeyTy &key) const {
+ return (uri == std::get<0>(key)) && (type == std::get<1>(key)) &&
+ (alignment == std::get<2>(key)) && (offset == std::get<3>(key)) &&
+ (size == std::get<4>(key)) && (nestedReferences == std::get<5>(key));
+ }
+
+ static llvm::hash_code hashKey(const KeyTy &key) {
+ return llvm::hash_combine(std::get<0>(key), std::get<1>(key),
+ std::get<2>(key), std::get<3>(key),
+ std::get<4>(key), std::get<5>(key));
+ }
+
+ static LoadedURIDenseResourceAttrStorage *
+ construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
+ auto uri = std::get<0>(key);
+ auto type = std::get<1>(key);
+ auto alignment = std::get<2>(key);
+ auto offset = std::get<3>(key);
+ auto size = std::get<4>(key);
+ auto nestedReferences = std::get<5>(key);
+ auto rawHandle = std::get<6>(key);
+ uri = allocator.copyInto(uri);
+ nestedReferences = allocator.copyInto(nestedReferences);
+
+ return new (allocator.allocate<LoadedURIDenseResourceAttrStorage>())
+ LoadedURIDenseResourceAttrStorage(uri, type, alignment, offset, size,
+ nestedReferences, rawHandle);
+ }
+
+ void initialize(MLIRContext *context) {
+ // If a handle is provided, then skip initializing.
+ if (rawHandle.has_value())
+ return;
+
+ // Init to failure. initialize is called after verify when initially
+ // created, so can't use absense of handle as failure.
+ rawHandle = failure();
+
+ // TODO: Handling different URI types.
+
+ StringRef fname = uri;
+ (void)fname.consume_front("file://");
+
+ std::unique_ptr<llvm::MemoryBuffer> buffer;
+ if (size.has_value() || offset.has_value()) {
+ std::optional<llvm::Align> mbAlign;
+ if (alignment)
+ mbAlign = llvm::Align(*alignment);
+ auto bufferOr = llvm::MemoryBuffer::getFileSlice(
+ fname, size.value_or(uint64_t(-1)), offset.value_or(0),
+ /*IsVolatile=*/false, mbAlign);
+ if (!bufferOr) {
+ // TODO: Find way to flag on location.
+ emitError(UnknownLoc::get(context))
+ << "loading URI failed: " << bufferOr.getError().message();
+ return;
+ }
+ buffer.swap(*bufferOr);
+ } else {
+ std::optional<llvm::Align> mbAlign;
+ if (alignment)
+ mbAlign = llvm::Align(*alignment);
+ auto bufferOr = llvm::MemoryBuffer::getFile(
+ fname, /*IsText=*/false, /*RequiresNullTerminator=*/true,
+ /*IsVolatile=*/false, mbAlign);
+ if (!bufferOr) {
+ // TODO: Find way to flag on location.
+ emitError(UnknownLoc::get(context))
+ << "loading URI failed: " << bufferOr.getError().message();
+ return;
+ }
+ buffer.swap(*bufferOr);
+ }
+
+ // TODO: Handling of different file types.
+
+ auto *data = buffer.release();
+ auto deleter = [data](void *, unsigned long, unsigned long) {
+ delete data;
+ };
+
+ // Extract the builtin dialect resource manager from context and
+ // construct a handle by inserting a new resource using the
+ // provided blob.
+ auto &manager =
+ DenseResourceElementsHandle::getManagerInterface(type.getContext());
+ auto str = data->getBuffer();
+ auto blob = UnmanagedAsmResourceBlob::allocateWithAlign(
+ {str.data(), str.size()}, alignment.value_or(32), deleter);
+ rawHandle = manager.insert("luri", std::move(blob));
+ }
+
+ StringRef uri;
+ ShapedType type;
+ bool isSplat = false;
+ ArrayRef<FlatSymbolRefAttr> nestedReferences;
+ std::optional<int64_t> alignment;
+ std::optional<int64_t> offset;
+ std::optional<int64_t> size;
+ std::optional<FailureOr<DenseResourceElementsHandle>> rawHandle;
+};
+} // namespace detail
+
+LoadedURIDenseResourceAttr LoadedURIDenseResourceAttr::get(
+ StringRef uri, Type type, std::optional<int64_t> alignment,
+ std::optional<int64_t> offset, std::optional<int64_t> size,
+ ArrayRef<FlatSymbolRefAttr> nestedReferences) {
+ auto *ctxt = type.getContext();
+ return Base::get(ctxt, uri, type, alignment, offset, size, nestedReferences,
+ std::nullopt);
+}
+
+LoadedURIDenseResourceAttr LoadedURIDenseResourceAttr::getChecked(
+ function_ref<::mlir::InFlightDiagnostic()> emitError, StringRef uri,
+ Type type, std::optional<int64_t> alignment, std::optional<int64_t> offset,
+ std::optional<int64_t> size, ArrayRef<FlatSymbolRefAttr> nestedReferences) {
+ auto *ctxt = type.getContext();
+ return Base::getChecked(emitError, ctxt, uri, type, alignment, offset, size,
+ nestedReferences, std::nullopt);
+}
+
+LogicalResult LoadedURIDenseResourceAttr::verify(
+ function_ref<::mlir::InFlightDiagnostic()> emitError, StringRef uri,
+ Type type, std::optional<int64_t> alignment, std::optional<int64_t> offset,
+ std::optional<int64_t> size, ArrayRef<FlatSymbolRefAttr> nestedReferences,
+ std::optional<FailureOr<DenseResourceElementsHandle>> handle) {
+ if (handle.has_value() && failed(handle.value()))
+ return emitError() << "URI failed to load";
+ return success();
+}
+
+StringRef LoadedURIDenseResourceAttr::getURI() const { return getImpl()->uri; }
+
+ShapedType LoadedURIDenseResourceAttr::getType() const {
+ return getImpl()->type;
+}
+
+std::optional<int64_t> LoadedURIDenseResourceAttr::getAlignment() const {
+ return getImpl()->alignment;
+}
+
+std::optional<int64_t> LoadedURIDenseResourceAttr::getByteOffset() const {
+ return getImpl()->offset;
+}
+
+std::optional<int64_t> LoadedURIDenseResourceAttr::getByteSize() const {
+ return getImpl()->size;
+}
+
+ArrayRef<FlatSymbolRefAttr>
+LoadedURIDenseResourceAttr::getNestedReferences() const {
+ return getImpl()->nestedReferences;
+}
+
+ArrayRef<char> LoadedURIDenseResourceAttr::getRawData() const {
+ assert(getImpl()->rawHandle.has_value() && succeeded(*getImpl()->rawHandle));
+ return getImpl()->rawHandle->value().getBlob()->getData();
+}
+
+} // namespace mlir
+
//===----------------------------------------------------------------------===//
// Attribute Utilities
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/IR/luri.mlir b/mlir/test/IR/luri.mlir
new file mode 100644
index 000000000000000..d4083c47b562a9f
--- /dev/null
+++ b/mlir/test/IR/luri.mlir
@@ -0,0 +1,11 @@
+// RUN: cd %S && mlir-opt --canonicalize %s | FileCheck %s
+
+func.func @foo() -> (tensor<10xf32>) {
+ %0 = arith.constant luri<"luri.raw", byte_offset = 80, byte_size = 80> : tensor<10xf32>
+ %1 = arith.constant luri<"luri.raw", byte_offset = 0, byte_size = 80> : tensor<10xf32>
+ %c = arith.constant dense<1.0e5> : tensor<10xf32>
+ ...
[truncated]
|
| @@ -1065,6 +1069,110 @@ ShapedType Parser::parseElementsLiteralType(Type type) { | |||
| return sType; | |||
| } | |||
|
|
|||
| /// Parse a sparse elements attribute. | |||
| ArrayRef<FlatSymbolRefAttr> nestedReferences = {}); | ||
|
|
||
| /// The set of data types that can be iterated by this attribute. | ||
| // TODO: Expand as needed; consider maybe reuse with DenseElementsAttr. |
There was a problem hiding this comment.
Yeah, I think these days, in practice this API is less useful than the getDataAs-style APIs for accessing dense elements/resources.
| std::vector<FlatSymbolRefAttr> nestedRefs; | ||
| while (consumeIf(Token::colon)) { | ||
| // Check for the '::' prefix. | ||
| if (parseToken(Token::colon, "ex[ected '::' for nested reference")) |
There was a problem hiding this comment.
| if (parseToken(Token::colon, "ex[ected '::' for nested reference")) | |
| if (parseToken(Token::colon, "expected '::' for nested reference")) |
| nestedReferences, rawHandle); | ||
| } | ||
|
|
||
| void initialize(MLIRContext *context) { |
There was a problem hiding this comment.
I didn't know there was an initialize hook! When does it get called?
There was a problem hiding this comment.
Yes, I was doing something a lot more convoluted until I found it :) That is a good question: so it happens post initializing the attribute storage in getWithTypeID, but it happens after the first verify call (I may have traced the wrong flow ... I'll double check) which is why I have to use a FailureOr.
|
|
||
| StringRef uri; | ||
| ShapedType type; | ||
| bool isSplat = false; |
There was a problem hiding this comment.
Yes, I was trying to reuse this with DenseElementsAttr ... which turned out to be a bit of a story. I'll remove.
| std::optional<int64_t> size, ArrayRef<FlatSymbolRefAttr> nestedReferences, | ||
| std::optional<FailureOr<DenseResourceElementsHandle>> handle) { | ||
| if (handle.has_value() && failed(handle.value())) | ||
| return emitError() << "URI failed to load"; |
There was a problem hiding this comment.
can you put the URI in the error message?
There was a problem hiding this comment.
It's not clear to me that this should be a builtin here: seems a bit quite specific?
Also the current implementation only support one scheme, the "URL" is implicitly a file path apparently, and it's eagerly loaded. Adding more scheme seems necessarily intrusive here (you have to modify the builtin dialect): there does not seem to be any dialect-customization available.
I am wondering if the value here isn't in the mechanic more than the attribute itself: that is we could just provide the facilities for dialect to have this behavior easily, without providing an actual attribute?
| /// An attribute that parses/prints as a identifier (file & symbol) while in | ||
| /// memory representing the loaded attribute. | ||
|
|
||
| class LoadedURIDenseResourceAttr |
There was a problem hiding this comment.
Can this be mostly defined in ODS?
There was a problem hiding this comment.
Yes, I was struggling with it a bit. Its been a while ... hence figured get this for discussion and refactor in mechanical way.
Not sure I agree, I see this as no more specific than DenseElementsAttr.
URI for now is implicitly a file path yes. This would later become the fallback behavior (well as there has to be some). So you'd switch on scheme of the URI (here currently only file is supported so no switching or checking) and as fallback you make some attempt (and there loading a file is probably the simplest). The plan here is to have a static registry method on the type (that's the TODO there). This would allow external groups to register a handler without changing the builtin dialect (e.g., So eagerly loaded I call out in the PR description as something that could be added. I am in between on it honestly. I figured one could do it lazily first time accessed (make it mutable attribute), that has the downside that your IR may pass verification multiple times before it actually checks if URI can be loaded and fails. I figured that we can adjust based on actual data.
Could you expand? Similar to the resource blob kind, I think >90% of the cases may be handled by this without dialect specialization. Common cases will probably be a splatted file of constants or MLIR Bytecode, the for the more application focussed formats (e.g., safetensor and the like) one may get lucky with the raw format even (many of them are made to be mmapable, we may need to add a layout attribute there ... but I haven't tried that). And so the dialect doesn't have much impact, its just an ElementsAttr where we can avoid opening the same file redundantly. |
There is history behind builtin, and it's pretty much admitted that we likely wouldn't have most of the things "builtin" if they were added in today's MLIR land. (we even looked into how to move tensor and memref to their own dialects!). So the existence of something in builtin isn't a reason in itself to add more :) Right now this looks like a nice convenient tool for specific cases (looks almost like it could be an "example" of using resources somehow), but too far from something "core enough" to be there. While landing code incrementally is preferred, I'd also want to understand the end picture more clearly (how we'll handle multiple schemes, etc.). Basically this raises more design questions to me than anything. |
True, but I see almost nothing dialect specific here. Even today when we landed ArrayAttr, it could have been anywhere, but it mostly made sense here in an independent manner. Now we could have a file or uri dialect :)
Registry could be as simple as a static member on the attribute (I mean vector<function_ref<std::optional<MemoryBuffer>(StringRef uri)> registry;and static LogicalResult LoadedURIElementsAttr::registerHandler(function_ref<std::optional<MemoryBuffer>(StringRef)> handler);which is a switch on the schema of the URI and returns a MemoryBuffer (now one go really complicated to try and enable streaming in data over RPCs ... but that'll make any access variable cost which could be surprising). And so one iterates over attempting different matchers, stop when there is a match and proceed further. I decided (well ... decided ...) against a map of schema here as I don't think we'll actually have that many and it also allows for cases where the same handler could handle multiple schemas. Now for the actual file decoding below, it just takes a MemoryBuffer, but there one would need a vector of file type handlers too. The main reason I went for attribute is the uniqueing behavior, its just a ElementsAttr like any other (this is the interface of interest here for me) and one can use it different ways on different ops, while the backing storage is (ideally) just memory mapped from file and roundtrips "just as name". I considered making this a resource type instead, that wasn't as obvious how to get all the parts, but could revisit. This was also on the builtin dialect though, so doesn't remove that concern. |
I think that ArrayAttr makes sense to be "builtin", but today, I would actually argue against even having an ArrayAttr. These days, uses of ArrayAttr tend to be tech debt (using an ArrayAttr of size 2 to represent LLVM complex constants, for example, is insanity). I think having a URI dialect makes sense, because then extensions can be registered and managed through the dialect itself, rather than building yet another extension system. |
I don't see it that way: I have to digest the rest of your post, I don't quite grasp completely the "registry on the attribute" part right now (and also why wouldn't |
It would be, I just didn't add any registry yet. Didn't want to add mechanism while there is only one supported thing.
But would the extension look any different? We don't have open attribute system here nor would we want to have everyone modify upstream for their use downstream. So wouldn't the registry on the dialect look exactly the same? It would just change where the vector/map is stored (I mean I don't mind a different dialect and this is a nicer storage place, I just don't see how its avoiding an extra registry). Adding ops for loading of different types and making the dialect extensible is an option, that uses existing registry mechanisms, and then one defines a folder per op. But then one needs a mechanism to ensure one doesn't load it multiple times (which probably ends up being a registry again, where here that registry is attribute key :) ). I see this as a relatively simple thing, its a constant not stored in the IR. Its one that allows you to reference to multiple external files in one flow which the existing tooling could now how to interact with (note: this is probably true for a non-loaded URI string attribute too, it just makes for a simple linkage without much innovation)., and potentially opens door for something like elide-elementsattr-if-larger to actually produce output that is exact while not obscuring human readability. |
Add attribute type that behaves like a regular ElementsAttr with the exception that it is printed and parsed by using a URI (and offset and size within) and references within URI. This allows for an external to context and IR constant. It is still a constant, so its not for parameter case nor an extern constant.
The offset and size allows for having a single file with multiple values inside (and for some custom formats that are mmap'able this ends up being sufficient if the layout matches). While the nested references allows for usage in formats that have an indexing table (e.g., MLIR bytecode). These differ as to when considered: the offset & size is used to dictate which part to load into a memory buffer, while the references are for once the file is loaded.
In effect this is an attribute whose value doesn't get printed and that is usable where the external file will be available. These are also usable in cases where ElementsAttr is supported, such as ml_program globals or the like where it is not immediately a constant --- such usage would allow for some final level of customization/downstream adaption to decide what to inline when (ideally the correct ops would be created from the start, either with URI, with loaded URI or constant with loaded URI, so that frontend dictates these). The other potential use is in the elide large attribute case, where it could be outlined to files now and you get the small IR while having exactly the same behaviour when passed through again (while being friendly on editor/human reader).
Note: this is a minimal initial version, itdoes not yet add support for different URI nor file types, nor lazy loading.