1//===- BytecodeReader.cpp - MLIR Bytecode Reader --------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir/Bytecode/BytecodeReader.h"
10#include "mlir/AsmParser/AsmParser.h"
11#include "mlir/Bytecode/BytecodeImplementation.h"
12#include "mlir/Bytecode/BytecodeOpInterface.h"
13#include "mlir/Bytecode/Encoding.h"
14#include "mlir/IR/BuiltinOps.h"
15#include "mlir/IR/Diagnostics.h"
16#include "mlir/IR/OpImplementation.h"
17#include "mlir/IR/Verifier.h"
18#include "mlir/IR/Visitors.h"
19#include "mlir/Support/LLVM.h"
20#include "llvm/ADT/ArrayRef.h"
21#include "llvm/ADT/ScopeExit.h"
22#include "llvm/ADT/StringExtras.h"
23#include "llvm/ADT/StringRef.h"
24#include "llvm/Support/Endian.h"
25#include "llvm/Support/MemoryBufferRef.h"
26#include "llvm/Support/SourceMgr.h"
27
28#include <cstddef>
29#include <list>
30#include <memory>
31#include <numeric>
32#include <optional>
33
34#define DEBUG_TYPE "mlir-bytecode-reader"
35
36using namespace mlir;
37
38/// Stringify the given section ID.
39static std::string toString(bytecode::Section::ID sectionID) {
40 switch (sectionID) {
41 case bytecode::Section::kString:
42 return "String (0)";
43 case bytecode::Section::kDialect:
44 return "Dialect (1)";
45 case bytecode::Section::kAttrType:
46 return "AttrType (2)";
47 case bytecode::Section::kAttrTypeOffset:
48 return "AttrTypeOffset (3)";
49 case bytecode::Section::kIR:
50 return "IR (4)";
51 case bytecode::Section::kResource:
52 return "Resource (5)";
53 case bytecode::Section::kResourceOffset:
54 return "ResourceOffset (6)";
55 case bytecode::Section::kDialectVersions:
56 return "DialectVersions (7)";
57 case bytecode::Section::kProperties:
58 return "Properties (8)";
59 default:
60 return ("Unknown (" + Twine(static_cast<unsigned>(sectionID)) + ")").str();
61 }
62}
63
64/// Returns true if the given top-level section ID is optional.
65static bool isSectionOptional(bytecode::Section::ID sectionID, int version) {
66 switch (sectionID) {
67 case bytecode::Section::kString:
68 case bytecode::Section::kDialect:
69 case bytecode::Section::kAttrType:
70 case bytecode::Section::kAttrTypeOffset:
71 case bytecode::Section::kIR:
72 return false;
73 case bytecode::Section::kResource:
74 case bytecode::Section::kResourceOffset:
75 case bytecode::Section::kDialectVersions:
76 return true;
77 case bytecode::Section::kProperties:
78 return version < bytecode::kNativePropertiesEncoding;
79 default:
80 llvm_unreachable("unknown section ID");
81 }
82}
83
84//===----------------------------------------------------------------------===//
85// EncodingReader
86//===----------------------------------------------------------------------===//
87
88namespace {
89class EncodingReader {
90public:
91 explicit EncodingReader(ArrayRef<uint8_t> contents, Location fileLoc)
92 : buffer(contents), dataIt(buffer.begin()), fileLoc(fileLoc) {}
93 explicit EncodingReader(StringRef contents, Location fileLoc)
94 : EncodingReader({reinterpret_cast<const uint8_t *>(contents.data()),
95 contents.size()},
96 fileLoc) {}
97
98 /// Returns true if the entire section has been read.
99 bool empty() const { return dataIt == buffer.end(); }
100
101 /// Returns the remaining size of the bytecode.
102 size_t size() const { return buffer.end() - dataIt; }
103
104 /// Align the current reader position to the specified alignment.
105 LogicalResult alignTo(unsigned alignment) {
106 if (!llvm::isPowerOf2_32(Value: alignment))
107 return emitError(args: "expected alignment to be a power-of-two");
108
109 auto isUnaligned = [&](const uint8_t *ptr) {
110 return ((uintptr_t)ptr & (alignment - 1)) != 0;
111 };
112
113 // Shift the reader position to the next alignment boundary.
114 while (isUnaligned(dataIt)) {
115 uint8_t padding;
116 if (failed(Result: parseByte(value&: padding)))
117 return failure();
118 if (padding != bytecode::kAlignmentByte) {
119 return emitError(args: "expected alignment byte (0xCB), but got: '0x" +
120 llvm::utohexstr(X: padding) + "'");
121 }
122 }
123
124 // Ensure the data iterator is now aligned. This case is unlikely because we
125 // *just* went through the effort to align the data iterator.
126 if (LLVM_UNLIKELY(isUnaligned(dataIt))) {
127 return emitError(args: "expected data iterator aligned to ", args&: alignment,
128 args: ", but got pointer: '0x" +
129 llvm::utohexstr(X: (uintptr_t)dataIt) + "'");
130 }
131
132 return success();
133 }
134
135 /// Emit an error using the given arguments.
136 template <typename... Args>
137 InFlightDiagnostic emitError(Args &&...args) const {
138 return ::emitError(loc: fileLoc).append(std::forward<Args>(args)...);
139 }
140 InFlightDiagnostic emitError() const { return ::emitError(loc: fileLoc); }
141
142 /// Parse a single byte from the stream.
143 template <typename T>
144 LogicalResult parseByte(T &value) {
145 if (empty())
146 return emitError(args: "attempting to parse a byte at the end of the bytecode");
147 value = static_cast<T>(*dataIt++);
148 return success();
149 }
150 /// Parse a range of bytes of 'length' into the given result.
151 LogicalResult parseBytes(size_t length, ArrayRef<uint8_t> &result) {
152 if (length > size()) {
153 return emitError(args: "attempting to parse ", args&: length, args: " bytes when only ",
154 args: size(), args: " remain");
155 }
156 result = {dataIt, length};
157 dataIt += length;
158 return success();
159 }
160 /// Parse a range of bytes of 'length' into the given result, which can be
161 /// assumed to be large enough to hold `length`.
162 LogicalResult parseBytes(size_t length, uint8_t *result) {
163 if (length > size()) {
164 return emitError(args: "attempting to parse ", args&: length, args: " bytes when only ",
165 args: size(), args: " remain");
166 }
167 memcpy(dest: result, src: dataIt, n: length);
168 dataIt += length;
169 return success();
170 }
171
172 /// Parse an aligned blob of data, where the alignment was encoded alongside
173 /// the data.
174 LogicalResult parseBlobAndAlignment(ArrayRef<uint8_t> &data,
175 uint64_t &alignment) {
176 uint64_t dataSize;
177 if (failed(Result: parseVarInt(result&: alignment)) || failed(Result: parseVarInt(result&: dataSize)) ||
178 failed(Result: alignTo(alignment)))
179 return failure();
180 return parseBytes(length: dataSize, result&: data);
181 }
182
183 /// Parse a variable length encoded integer from the byte stream. The first
184 /// encoded byte contains a prefix in the low bits indicating the encoded
185 /// length of the value. This length prefix is a bit sequence of '0's followed
186 /// by a '1'. The number of '0' bits indicate the number of _additional_ bytes
187 /// (not including the prefix byte). All remaining bits in the first byte,
188 /// along with all of the bits in additional bytes, provide the value of the
189 /// integer encoded in little-endian order.
190 LogicalResult parseVarInt(uint64_t &result) {
191 // Parse the first byte of the encoding, which contains the length prefix.
192 if (failed(Result: parseByte(value&: result)))
193 return failure();
194
195 // Handle the overwhelmingly common case where the value is stored in a
196 // single byte. In this case, the first bit is the `1` marker bit.
197 if (LLVM_LIKELY(result & 1)) {
198 result >>= 1;
199 return success();
200 }
201
202 // Handle the overwhelming uncommon case where the value required all 8
203 // bytes (i.e. a really really big number). In this case, the marker byte is
204 // all zeros: `00000000`.
205 if (LLVM_UNLIKELY(result == 0)) {
206 llvm::support::ulittle64_t resultLE;
207 if (failed(Result: parseBytes(length: sizeof(resultLE),
208 result: reinterpret_cast<uint8_t *>(&resultLE))))
209 return failure();
210 result = resultLE;
211 return success();
212 }
213 return parseMultiByteVarInt(result);
214 }
215
216 /// Parse a signed variable length encoded integer from the byte stream. A
217 /// signed varint is encoded as a normal varint with zigzag encoding applied,
218 /// i.e. the low bit of the value is used to indicate the sign.
219 LogicalResult parseSignedVarInt(uint64_t &result) {
220 if (failed(Result: parseVarInt(result)))
221 return failure();
222 // Essentially (but using unsigned): (x >> 1) ^ -(x & 1)
223 result = (result >> 1) ^ (~(result & 1) + 1);
224 return success();
225 }
226
227 /// Parse a variable length encoded integer whose low bit is used to encode an
228 /// unrelated flag, i.e: `(integerValue << 1) | (flag ? 1 : 0)`.
229 LogicalResult parseVarIntWithFlag(uint64_t &result, bool &flag) {
230 if (failed(Result: parseVarInt(result)))
231 return failure();
232 flag = result & 1;
233 result >>= 1;
234 return success();
235 }
236
237 /// Skip the first `length` bytes within the reader.
238 LogicalResult skipBytes(size_t length) {
239 if (length > size()) {
240 return emitError(args: "attempting to skip ", args&: length, args: " bytes when only ",
241 args: size(), args: " remain");
242 }
243 dataIt += length;
244 return success();
245 }
246
247 /// Parse a null-terminated string into `result` (without including the NUL
248 /// terminator).
249 LogicalResult parseNullTerminatedString(StringRef &result) {
250 const char *startIt = (const char *)dataIt;
251 const char *nulIt = (const char *)memchr(s: startIt, c: 0, n: size());
252 if (!nulIt)
253 return emitError(
254 args: "malformed null-terminated string, no null character found");
255
256 result = StringRef(startIt, nulIt - startIt);
257 dataIt = (const uint8_t *)nulIt + 1;
258 return success();
259 }
260
261 /// Parse a section header, placing the kind of section in `sectionID` and the
262 /// contents of the section in `sectionData`.
263 LogicalResult parseSection(bytecode::Section::ID &sectionID,
264 ArrayRef<uint8_t> &sectionData) {
265 uint8_t sectionIDAndHasAlignment;
266 uint64_t length;
267 if (failed(Result: parseByte(value&: sectionIDAndHasAlignment)) ||
268 failed(Result: parseVarInt(result&: length)))
269 return failure();
270
271 // Extract the section ID and whether the section is aligned. The high bit
272 // of the ID is the alignment flag.
273 sectionID = static_cast<bytecode::Section::ID>(sectionIDAndHasAlignment &
274 0b01111111);
275 bool hasAlignment = sectionIDAndHasAlignment & 0b10000000;
276
277 // Check that the section is actually valid before trying to process its
278 // data.
279 if (sectionID >= bytecode::Section::kNumSections)
280 return emitError(args: "invalid section ID: ", args: unsigned(sectionID));
281
282 // Process the section alignment if present.
283 if (hasAlignment) {
284 uint64_t alignment;
285 if (failed(Result: parseVarInt(result&: alignment)) || failed(Result: alignTo(alignment)))
286 return failure();
287 }
288
289 // Parse the actual section data.
290 return parseBytes(length: static_cast<size_t>(length), result&: sectionData);
291 }
292
293 Location getLoc() const { return fileLoc; }
294
295private:
296 /// Parse a variable length encoded integer from the byte stream. This method
297 /// is a fallback when the number of bytes used to encode the value is greater
298 /// than 1, but less than the max (9). The provided `result` value can be
299 /// assumed to already contain the first byte of the value.
300 /// NOTE: This method is marked noinline to avoid pessimizing the common case
301 /// of single byte encoding.
302 LLVM_ATTRIBUTE_NOINLINE LogicalResult parseMultiByteVarInt(uint64_t &result) {
303 // Count the number of trailing zeros in the marker byte, this indicates the
304 // number of trailing bytes that are part of the value. We use `uint32_t`
305 // here because we only care about the first byte, and so that be actually
306 // get ctz intrinsic calls when possible (the `uint8_t` overload uses a loop
307 // implementation).
308 uint32_t numBytes = llvm::countr_zero<uint32_t>(Val: result);
309 assert(numBytes > 0 && numBytes <= 7 &&
310 "unexpected number of trailing zeros in varint encoding");
311
312 // Parse in the remaining bytes of the value.
313 llvm::support::ulittle64_t resultLE(result);
314 if (failed(
315 Result: parseBytes(length: numBytes, result: reinterpret_cast<uint8_t *>(&resultLE) + 1)))
316 return failure();
317
318 // Shift out the low-order bits that were used to mark how the value was
319 // encoded.
320 result = resultLE >> (numBytes + 1);
321 return success();
322 }
323
324 /// The bytecode buffer.
325 ArrayRef<uint8_t> buffer;
326
327 /// The current iterator within the 'buffer'.
328 const uint8_t *dataIt;
329
330 /// A location for the bytecode used to report errors.
331 Location fileLoc;
332};
333} // namespace
334
335/// Resolve an index into the given entry list. `entry` may either be a
336/// reference, in which case it is assigned to the corresponding value in
337/// `entries`, or a pointer, in which case it is assigned to the address of the
338/// element in `entries`.
339template <typename RangeT, typename T>
340static LogicalResult resolveEntry(EncodingReader &reader, RangeT &entries,
341 uint64_t index, T &entry,
342 StringRef entryStr) {
343 if (index >= entries.size())
344 return reader.emitError(args: "invalid ", args&: entryStr, args: " index: ", args&: index);
345
346 // If the provided entry is a pointer, resolve to the address of the entry.
347 if constexpr (std::is_convertible_v<llvm::detail::ValueOfRange<RangeT>, T>)
348 entry = entries[index];
349 else
350 entry = &entries[index];
351 return success();
352}
353
354/// Parse and resolve an index into the given entry list.
355template <typename RangeT, typename T>
356static LogicalResult parseEntry(EncodingReader &reader, RangeT &entries,
357 T &entry, StringRef entryStr) {
358 uint64_t entryIdx;
359 if (failed(Result: reader.parseVarInt(result&: entryIdx)))
360 return failure();
361 return resolveEntry(reader, entries, entryIdx, entry, entryStr);
362}
363
364//===----------------------------------------------------------------------===//
365// StringSectionReader
366//===----------------------------------------------------------------------===//
367
368namespace {
369/// This class is used to read references to the string section from the
370/// bytecode.
371class StringSectionReader {
372public:
373 /// Initialize the string section reader with the given section data.
374 LogicalResult initialize(Location fileLoc, ArrayRef<uint8_t> sectionData);
375
376 /// Parse a shared string from the string section. The shared string is
377 /// encoded using an index to a corresponding string in the string section.
378 LogicalResult parseString(EncodingReader &reader, StringRef &result) const {
379 return parseEntry(reader, entries: strings, entry&: result, entryStr: "string");
380 }
381
382 /// Parse a shared string from the string section. The shared string is
383 /// encoded using an index to a corresponding string in the string section.
384 /// This variant parses a flag compressed with the index.
385 LogicalResult parseStringWithFlag(EncodingReader &reader, StringRef &result,
386 bool &flag) const {
387 uint64_t entryIdx;
388 if (failed(Result: reader.parseVarIntWithFlag(result&: entryIdx, flag)))
389 return failure();
390 return parseStringAtIndex(reader, index: entryIdx, result);
391 }
392
393 /// Parse a shared string from the string section. The shared string is
394 /// encoded using an index to a corresponding string in the string section.
395 LogicalResult parseStringAtIndex(EncodingReader &reader, uint64_t index,
396 StringRef &result) const {
397 return resolveEntry(reader, entries: strings, index, entry&: result, entryStr: "string");
398 }
399
400private:
401 /// The table of strings referenced within the bytecode file.
402 SmallVector<StringRef> strings;
403};
404} // namespace
405
406LogicalResult StringSectionReader::initialize(Location fileLoc,
407 ArrayRef<uint8_t> sectionData) {
408 EncodingReader stringReader(sectionData, fileLoc);
409
410 // Parse the number of strings in the section.
411 uint64_t numStrings;
412 if (failed(Result: stringReader.parseVarInt(result&: numStrings)))
413 return failure();
414 strings.resize(N: numStrings);
415
416 // Parse each of the strings. The sizes of the strings are encoded in reverse
417 // order, so that's the order we populate the table.
418 size_t stringDataEndOffset = sectionData.size();
419 for (StringRef &string : llvm::reverse(C&: strings)) {
420 uint64_t stringSize;
421 if (failed(Result: stringReader.parseVarInt(result&: stringSize)))
422 return failure();
423 if (stringDataEndOffset < stringSize) {
424 return stringReader.emitError(
425 args: "string size exceeds the available data size");
426 }
427
428 // Extract the string from the data, dropping the null character.
429 size_t stringOffset = stringDataEndOffset - stringSize;
430 string = StringRef(
431 reinterpret_cast<const char *>(sectionData.data() + stringOffset),
432 stringSize - 1);
433 stringDataEndOffset = stringOffset;
434 }
435
436 // Check that the only remaining data was for the strings, i.e. the reader
437 // should be at the same offset as the first string.
438 if ((sectionData.size() - stringReader.size()) != stringDataEndOffset) {
439 return stringReader.emitError(args: "unexpected trailing data between the "
440 "offsets for strings and their data");
441 }
442 return success();
443}
444
445//===----------------------------------------------------------------------===//
446// BytecodeDialect
447//===----------------------------------------------------------------------===//
448
449namespace {
450class DialectReader;
451
452/// This struct represents a dialect entry within the bytecode.
453struct BytecodeDialect {
454 /// Load the dialect into the provided context if it hasn't been loaded yet.
455 /// Returns failure if the dialect couldn't be loaded *and* the provided
456 /// context does not allow unregistered dialects. The provided reader is used
457 /// for error emission if necessary.
458 LogicalResult load(const DialectReader &reader, MLIRContext *ctx);
459
460 /// Return the loaded dialect, or nullptr if the dialect is unknown. This can
461 /// only be called after `load`.
462 Dialect *getLoadedDialect() const {
463 assert(dialect &&
464 "expected `load` to be invoked before `getLoadedDialect`");
465 return *dialect;
466 }
467
468 /// The loaded dialect entry. This field is std::nullopt if we haven't
469 /// attempted to load, nullptr if we failed to load, otherwise the loaded
470 /// dialect.
471 std::optional<Dialect *> dialect;
472
473 /// The bytecode interface of the dialect, or nullptr if the dialect does not
474 /// implement the bytecode interface. This field should only be checked if the
475 /// `dialect` field is not std::nullopt.
476 const BytecodeDialectInterface *interface = nullptr;
477
478 /// The name of the dialect.
479 StringRef name;
480
481 /// A buffer containing the encoding of the dialect version parsed.
482 ArrayRef<uint8_t> versionBuffer;
483
484 /// Lazy loaded dialect version from the handle above.
485 std::unique_ptr<DialectVersion> loadedVersion;
486};
487
488/// This struct represents an operation name entry within the bytecode.
489struct BytecodeOperationName {
490 BytecodeOperationName(BytecodeDialect *dialect, StringRef name,
491 std::optional<bool> wasRegistered)
492 : dialect(dialect), name(name), wasRegistered(wasRegistered) {}
493
494 /// The loaded operation name, or std::nullopt if it hasn't been processed
495 /// yet.
496 std::optional<OperationName> opName;
497
498 /// The dialect that owns this operation name.
499 BytecodeDialect *dialect;
500
501 /// The name of the operation, without the dialect prefix.
502 StringRef name;
503
504 /// Whether this operation was registered when the bytecode was produced.
505 /// This flag is populated when bytecode version >=kNativePropertiesEncoding.
506 std::optional<bool> wasRegistered;
507};
508} // namespace
509
510/// Parse a single dialect group encoded in the byte stream.
511static LogicalResult parseDialectGrouping(
512 EncodingReader &reader,
513 MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
514 function_ref<LogicalResult(BytecodeDialect *)> entryCallback) {
515 // Parse the dialect and the number of entries in the group.
516 std::unique_ptr<BytecodeDialect> *dialect;
517 if (failed(Result: parseEntry(reader, entries&: dialects, entry&: dialect, entryStr: "dialect")))
518 return failure();
519 uint64_t numEntries;
520 if (failed(Result: reader.parseVarInt(result&: numEntries)))
521 return failure();
522
523 for (uint64_t i = 0; i < numEntries; ++i)
524 if (failed(Result: entryCallback(dialect->get())))
525 return failure();
526 return success();
527}
528
529//===----------------------------------------------------------------------===//
530// ResourceSectionReader
531//===----------------------------------------------------------------------===//
532
533namespace {
534/// This class is used to read the resource section from the bytecode.
535class ResourceSectionReader {
536public:
537 /// Initialize the resource section reader with the given section data.
538 LogicalResult
539 initialize(Location fileLoc, const ParserConfig &config,
540 MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
541 StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
542 ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
543 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef);
544
545 /// Parse a dialect resource handle from the resource section.
546 LogicalResult parseResourceHandle(EncodingReader &reader,
547 AsmDialectResourceHandle &result) const {
548 return parseEntry(reader, entries: dialectResources, entry&: result, entryStr: "resource handle");
549 }
550
551private:
552 /// The table of dialect resources within the bytecode file.
553 SmallVector<AsmDialectResourceHandle> dialectResources;
554 llvm::StringMap<std::string> dialectResourceHandleRenamingMap;
555};
556
557class ParsedResourceEntry : public AsmParsedResourceEntry {
558public:
559 ParsedResourceEntry(StringRef key, AsmResourceEntryKind kind,
560 EncodingReader &reader, StringSectionReader &stringReader,
561 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
562 : key(key), kind(kind), reader(reader), stringReader(stringReader),
563 bufferOwnerRef(bufferOwnerRef) {}
564 ~ParsedResourceEntry() override = default;
565
566 StringRef getKey() const final { return key; }
567
568 InFlightDiagnostic emitError() const final { return reader.emitError(); }
569
570 AsmResourceEntryKind getKind() const final { return kind; }
571
572 FailureOr<bool> parseAsBool() const final {
573 if (kind != AsmResourceEntryKind::Bool)
574 return emitError() << "expected a bool resource entry, but found a "
575 << toString(kind) << " entry instead";
576
577 bool value;
578 if (failed(Result: reader.parseByte(value)))
579 return failure();
580 return value;
581 }
582 FailureOr<std::string> parseAsString() const final {
583 if (kind != AsmResourceEntryKind::String)
584 return emitError() << "expected a string resource entry, but found a "
585 << toString(kind) << " entry instead";
586
587 StringRef string;
588 if (failed(Result: stringReader.parseString(reader, result&: string)))
589 return failure();
590 return string.str();
591 }
592
593 FailureOr<AsmResourceBlob>
594 parseAsBlob(BlobAllocatorFn allocator) const final {
595 if (kind != AsmResourceEntryKind::Blob)
596 return emitError() << "expected a blob resource entry, but found a "
597 << toString(kind) << " entry instead";
598
599 ArrayRef<uint8_t> data;
600 uint64_t alignment;
601 if (failed(Result: reader.parseBlobAndAlignment(data, alignment)))
602 return failure();
603
604 // If we have an extendable reference to the buffer owner, we don't need to
605 // allocate a new buffer for the data, and can use the data directly.
606 if (bufferOwnerRef) {
607 ArrayRef<char> charData(reinterpret_cast<const char *>(data.data()),
608 data.size());
609
610 // Allocate an unmanager buffer which captures a reference to the owner.
611 // For now we just mark this as immutable, but in the future we should
612 // explore marking this as mutable when desired.
613 return UnmanagedAsmResourceBlob::allocateWithAlign(
614 data: charData, align: alignment,
615 deleter: [bufferOwnerRef = bufferOwnerRef](void *, size_t, size_t) {});
616 }
617
618 // Allocate memory for the blob using the provided allocator and copy the
619 // data into it.
620 AsmResourceBlob blob = allocator(data.size(), alignment);
621 assert(llvm::isAddrAligned(llvm::Align(alignment), blob.getData().data()) &&
622 blob.isMutable() &&
623 "blob allocator did not return a properly aligned address");
624 memcpy(dest: blob.getMutableData().data(), src: data.data(), n: data.size());
625 return blob;
626 }
627
628private:
629 StringRef key;
630 AsmResourceEntryKind kind;
631 EncodingReader &reader;
632 StringSectionReader &stringReader;
633 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
634};
635} // namespace
636
637template <typename T>
638static LogicalResult
639parseResourceGroup(Location fileLoc, bool allowEmpty,
640 EncodingReader &offsetReader, EncodingReader &resourceReader,
641 StringSectionReader &stringReader, T *handler,
642 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef,
643 function_ref<StringRef(StringRef)> remapKey = {},
644 function_ref<LogicalResult(StringRef)> processKeyFn = {}) {
645 uint64_t numResources;
646 if (failed(Result: offsetReader.parseVarInt(result&: numResources)))
647 return failure();
648
649 for (uint64_t i = 0; i < numResources; ++i) {
650 StringRef key;
651 AsmResourceEntryKind kind;
652 uint64_t resourceOffset;
653 ArrayRef<uint8_t> data;
654 if (failed(Result: stringReader.parseString(reader&: offsetReader, result&: key)) ||
655 failed(Result: offsetReader.parseVarInt(result&: resourceOffset)) ||
656 failed(Result: offsetReader.parseByte(value&: kind)) ||
657 failed(Result: resourceReader.parseBytes(length: resourceOffset, result&: data)))
658 return failure();
659
660 // Process the resource key.
661 if ((processKeyFn && failed(Result: processKeyFn(key))))
662 return failure();
663
664 // If the resource data is empty and we allow it, don't error out when
665 // parsing below, just skip it.
666 if (allowEmpty && data.empty())
667 continue;
668
669 // Ignore the entry if we don't have a valid handler.
670 if (!handler)
671 continue;
672
673 // Otherwise, parse the resource value.
674 EncodingReader entryReader(data, fileLoc);
675 key = remapKey(key);
676 ParsedResourceEntry entry(key, kind, entryReader, stringReader,
677 bufferOwnerRef);
678 if (failed(handler->parseResource(entry)))
679 return failure();
680 if (!entryReader.empty()) {
681 return entryReader.emitError(
682 args: "unexpected trailing bytes in resource entry '", args&: key, args: "'");
683 }
684 }
685 return success();
686}
687
688LogicalResult ResourceSectionReader::initialize(
689 Location fileLoc, const ParserConfig &config,
690 MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
691 StringSectionReader &stringReader, ArrayRef<uint8_t> sectionData,
692 ArrayRef<uint8_t> offsetSectionData, DialectReader &dialectReader,
693 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
694 EncodingReader resourceReader(sectionData, fileLoc);
695 EncodingReader offsetReader(offsetSectionData, fileLoc);
696
697 // Read the number of external resource providers.
698 uint64_t numExternalResourceGroups;
699 if (failed(Result: offsetReader.parseVarInt(result&: numExternalResourceGroups)))
700 return failure();
701
702 // Utility functor that dispatches to `parseResourceGroup`, but implicitly
703 // provides most of the arguments.
704 auto parseGroup = [&](auto *handler, bool allowEmpty = false,
705 function_ref<LogicalResult(StringRef)> keyFn = {}) {
706 auto resolveKey = [&](StringRef key) -> StringRef {
707 auto it = dialectResourceHandleRenamingMap.find(Key: key);
708 if (it == dialectResourceHandleRenamingMap.end())
709 return key;
710 return it->second;
711 };
712
713 return parseResourceGroup(fileLoc, allowEmpty, offsetReader, resourceReader,
714 stringReader, handler, bufferOwnerRef, resolveKey,
715 keyFn);
716 };
717
718 // Read the external resources from the bytecode.
719 for (uint64_t i = 0; i < numExternalResourceGroups; ++i) {
720 StringRef key;
721 if (failed(Result: stringReader.parseString(reader&: offsetReader, result&: key)))
722 return failure();
723
724 // Get the handler for these resources.
725 // TODO: Should we require handling external resources in some scenarios?
726 AsmResourceParser *handler = config.getResourceParser(name: key);
727 if (!handler) {
728 emitWarning(loc: fileLoc) << "ignoring unknown external resources for '" << key
729 << "'";
730 }
731
732 if (failed(Result: parseGroup(handler)))
733 return failure();
734 }
735
736 // Read the dialect resources from the bytecode.
737 MLIRContext *ctx = fileLoc->getContext();
738 while (!offsetReader.empty()) {
739 std::unique_ptr<BytecodeDialect> *dialect;
740 if (failed(Result: parseEntry(reader&: offsetReader, entries&: dialects, entry&: dialect, entryStr: "dialect")) ||
741 failed(Result: (*dialect)->load(reader: dialectReader, ctx)))
742 return failure();
743 Dialect *loadedDialect = (*dialect)->getLoadedDialect();
744 if (!loadedDialect) {
745 return resourceReader.emitError()
746 << "dialect '" << (*dialect)->name << "' is unknown";
747 }
748 const auto *handler = dyn_cast<OpAsmDialectInterface>(Val: loadedDialect);
749 if (!handler) {
750 return resourceReader.emitError()
751 << "unexpected resources for dialect '" << (*dialect)->name << "'";
752 }
753
754 // Ensure that each resource is declared before being processed.
755 auto processResourceKeyFn = [&](StringRef key) -> LogicalResult {
756 FailureOr<AsmDialectResourceHandle> handle =
757 handler->declareResource(key);
758 if (failed(Result: handle)) {
759 return resourceReader.emitError()
760 << "unknown 'resource' key '" << key << "' for dialect '"
761 << (*dialect)->name << "'";
762 }
763 dialectResourceHandleRenamingMap[key] = handler->getResourceKey(handle: *handle);
764 dialectResources.push_back(Elt: *handle);
765 return success();
766 };
767
768 // Parse the resources for this dialect. We allow empty resources because we
769 // just treat these as declarations.
770 if (failed(Result: parseGroup(handler, /*allowEmpty=*/true, processResourceKeyFn)))
771 return failure();
772 }
773
774 return success();
775}
776
777//===----------------------------------------------------------------------===//
778// Attribute/Type Reader
779//===----------------------------------------------------------------------===//
780
781namespace {
782/// This class provides support for reading attribute and type entries from the
783/// bytecode. Attribute and Type entries are read lazily on demand, so we use
784/// this reader to manage when to actually parse them from the bytecode.
785class AttrTypeReader {
786 /// This class represents a single attribute or type entry.
787 template <typename T>
788 struct Entry {
789 /// The entry, or null if it hasn't been resolved yet.
790 T entry = {};
791 /// The parent dialect of this entry.
792 BytecodeDialect *dialect = nullptr;
793 /// A flag indicating if the entry was encoded using a custom encoding,
794 /// instead of using the textual assembly format.
795 bool hasCustomEncoding = false;
796 /// The raw data of this entry in the bytecode.
797 ArrayRef<uint8_t> data;
798 };
799 using AttrEntry = Entry<Attribute>;
800 using TypeEntry = Entry<Type>;
801
802public:
803 AttrTypeReader(const StringSectionReader &stringReader,
804 const ResourceSectionReader &resourceReader,
805 const llvm::StringMap<BytecodeDialect *> &dialectsMap,
806 uint64_t &bytecodeVersion, Location fileLoc,
807 const ParserConfig &config)
808 : stringReader(stringReader), resourceReader(resourceReader),
809 dialectsMap(dialectsMap), fileLoc(fileLoc),
810 bytecodeVersion(bytecodeVersion), parserConfig(config) {}
811
812 /// Initialize the attribute and type information within the reader.
813 LogicalResult
814 initialize(MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
815 ArrayRef<uint8_t> sectionData,
816 ArrayRef<uint8_t> offsetSectionData);
817
818 /// Resolve the attribute or type at the given index. Returns nullptr on
819 /// failure.
820 Attribute resolveAttribute(size_t index) {
821 return resolveEntry(entries&: attributes, index, entryType: "Attribute");
822 }
823 Type resolveType(size_t index) { return resolveEntry(entries&: types, index, entryType: "Type"); }
824
825 /// Parse a reference to an attribute or type using the given reader.
826 LogicalResult parseAttribute(EncodingReader &reader, Attribute &result) {
827 uint64_t attrIdx;
828 if (failed(Result: reader.parseVarInt(result&: attrIdx)))
829 return failure();
830 result = resolveAttribute(index: attrIdx);
831 return success(IsSuccess: !!result);
832 }
833 LogicalResult parseOptionalAttribute(EncodingReader &reader,
834 Attribute &result) {
835 uint64_t attrIdx;
836 bool flag;
837 if (failed(Result: reader.parseVarIntWithFlag(result&: attrIdx, flag)))
838 return failure();
839 if (!flag)
840 return success();
841 result = resolveAttribute(index: attrIdx);
842 return success(IsSuccess: !!result);
843 }
844
845 LogicalResult parseType(EncodingReader &reader, Type &result) {
846 uint64_t typeIdx;
847 if (failed(Result: reader.parseVarInt(result&: typeIdx)))
848 return failure();
849 result = resolveType(index: typeIdx);
850 return success(IsSuccess: !!result);
851 }
852
853 template <typename T>
854 LogicalResult parseAttribute(EncodingReader &reader, T &result) {
855 Attribute baseResult;
856 if (failed(Result: parseAttribute(reader, result&: baseResult)))
857 return failure();
858 if ((result = dyn_cast<T>(baseResult)))
859 return success();
860 return reader.emitError("expected attribute of type: ",
861 llvm::getTypeName<T>(), ", but got: ", baseResult);
862 }
863
864private:
865 /// Resolve the given entry at `index`.
866 template <typename T>
867 T resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
868 StringRef entryType);
869
870 /// Parse an entry using the given reader that was encoded using the textual
871 /// assembly format.
872 template <typename T>
873 LogicalResult parseAsmEntry(T &result, EncodingReader &reader,
874 StringRef entryType);
875
876 /// Parse an entry using the given reader that was encoded using a custom
877 /// bytecode format.
878 template <typename T>
879 LogicalResult parseCustomEntry(Entry<T> &entry, EncodingReader &reader,
880 StringRef entryType);
881
882 /// The string section reader used to resolve string references when parsing
883 /// custom encoded attribute/type entries.
884 const StringSectionReader &stringReader;
885
886 /// The resource section reader used to resolve resource references when
887 /// parsing custom encoded attribute/type entries.
888 const ResourceSectionReader &resourceReader;
889
890 /// The map of the loaded dialects used to retrieve dialect information, such
891 /// as the dialect version.
892 const llvm::StringMap<BytecodeDialect *> &dialectsMap;
893
894 /// The set of attribute and type entries.
895 SmallVector<AttrEntry> attributes;
896 SmallVector<TypeEntry> types;
897
898 /// A location used for error emission.
899 Location fileLoc;
900
901 /// Current bytecode version being used.
902 uint64_t &bytecodeVersion;
903
904 /// Reference to the parser configuration.
905 const ParserConfig &parserConfig;
906};
907
908class DialectReader : public DialectBytecodeReader {
909public:
910 DialectReader(AttrTypeReader &attrTypeReader,
911 const StringSectionReader &stringReader,
912 const ResourceSectionReader &resourceReader,
913 const llvm::StringMap<BytecodeDialect *> &dialectsMap,
914 EncodingReader &reader, uint64_t &bytecodeVersion)
915 : attrTypeReader(attrTypeReader), stringReader(stringReader),
916 resourceReader(resourceReader), dialectsMap(dialectsMap),
917 reader(reader), bytecodeVersion(bytecodeVersion) {}
918
919 InFlightDiagnostic emitError(const Twine &msg) const override {
920 return reader.emitError(args: msg);
921 }
922
923 FailureOr<const DialectVersion *>
924 getDialectVersion(StringRef dialectName) const override {
925 // First check if the dialect is available in the map.
926 auto dialectEntry = dialectsMap.find(Key: dialectName);
927 if (dialectEntry == dialectsMap.end())
928 return failure();
929 // If the dialect was found, try to load it. This will trigger reading the
930 // bytecode version from the version buffer if it wasn't already processed.
931 // Return failure if either of those two actions could not be completed.
932 if (failed(Result: dialectEntry->getValue()->load(reader: *this, ctx: getLoc().getContext())) ||
933 dialectEntry->getValue()->loadedVersion == nullptr)
934 return failure();
935 return dialectEntry->getValue()->loadedVersion.get();
936 }
937
938 MLIRContext *getContext() const override { return getLoc().getContext(); }
939
940 uint64_t getBytecodeVersion() const override { return bytecodeVersion; }
941
942 DialectReader withEncodingReader(EncodingReader &encReader) const {
943 return DialectReader(attrTypeReader, stringReader, resourceReader,
944 dialectsMap, encReader, bytecodeVersion);
945 }
946
947 Location getLoc() const { return reader.getLoc(); }
948
949 //===--------------------------------------------------------------------===//
950 // IR
951 //===--------------------------------------------------------------------===//
952
953 LogicalResult readAttribute(Attribute &result) override {
954 return attrTypeReader.parseAttribute(reader, result);
955 }
956 LogicalResult readOptionalAttribute(Attribute &result) override {
957 return attrTypeReader.parseOptionalAttribute(reader, result);
958 }
959 LogicalResult readType(Type &result) override {
960 return attrTypeReader.parseType(reader, result);
961 }
962
963 FailureOr<AsmDialectResourceHandle> readResourceHandle() override {
964 AsmDialectResourceHandle handle;
965 if (failed(Result: resourceReader.parseResourceHandle(reader, result&: handle)))
966 return failure();
967 return handle;
968 }
969
970 //===--------------------------------------------------------------------===//
971 // Primitives
972 //===--------------------------------------------------------------------===//
973
974 LogicalResult readVarInt(uint64_t &result) override {
975 return reader.parseVarInt(result);
976 }
977
978 LogicalResult readSignedVarInt(int64_t &result) override {
979 uint64_t unsignedResult;
980 if (failed(Result: reader.parseSignedVarInt(result&: unsignedResult)))
981 return failure();
982 result = static_cast<int64_t>(unsignedResult);
983 return success();
984 }
985
986 FailureOr<APInt> readAPIntWithKnownWidth(unsigned bitWidth) override {
987 // Small values are encoded using a single byte.
988 if (bitWidth <= 8) {
989 uint8_t value;
990 if (failed(Result: reader.parseByte(value)))
991 return failure();
992 return APInt(bitWidth, value);
993 }
994
995 // Large values up to 64 bits are encoded using a single varint.
996 if (bitWidth <= 64) {
997 uint64_t value;
998 if (failed(Result: reader.parseSignedVarInt(result&: value)))
999 return failure();
1000 return APInt(bitWidth, value);
1001 }
1002
1003 // Otherwise, for really big values we encode the array of active words in
1004 // the value.
1005 uint64_t numActiveWords;
1006 if (failed(Result: reader.parseVarInt(result&: numActiveWords)))
1007 return failure();
1008 SmallVector<uint64_t, 4> words(numActiveWords);
1009 for (uint64_t i = 0; i < numActiveWords; ++i)
1010 if (failed(Result: reader.parseSignedVarInt(result&: words[i])))
1011 return failure();
1012 return APInt(bitWidth, words);
1013 }
1014
1015 FailureOr<APFloat>
1016 readAPFloatWithKnownSemantics(const llvm::fltSemantics &semantics) override {
1017 FailureOr<APInt> intVal =
1018 readAPIntWithKnownWidth(bitWidth: APFloat::getSizeInBits(Sem: semantics));
1019 if (failed(Result: intVal))
1020 return failure();
1021 return APFloat(semantics, *intVal);
1022 }
1023
1024 LogicalResult readString(StringRef &result) override {
1025 return stringReader.parseString(reader, result);
1026 }
1027
1028 LogicalResult readBlob(ArrayRef<char> &result) override {
1029 uint64_t dataSize;
1030 ArrayRef<uint8_t> data;
1031 if (failed(Result: reader.parseVarInt(result&: dataSize)) ||
1032 failed(Result: reader.parseBytes(length: dataSize, result&: data)))
1033 return failure();
1034 result = llvm::ArrayRef(reinterpret_cast<const char *>(data.data()),
1035 data.size());
1036 return success();
1037 }
1038
1039 LogicalResult readBool(bool &result) override {
1040 return reader.parseByte(value&: result);
1041 }
1042
1043private:
1044 AttrTypeReader &attrTypeReader;
1045 const StringSectionReader &stringReader;
1046 const ResourceSectionReader &resourceReader;
1047 const llvm::StringMap<BytecodeDialect *> &dialectsMap;
1048 EncodingReader &reader;
1049 uint64_t &bytecodeVersion;
1050};
1051
1052/// Wraps the properties section and handles reading properties out of it.
1053class PropertiesSectionReader {
1054public:
1055 /// Initialize the properties section reader with the given section data.
1056 LogicalResult initialize(Location fileLoc, ArrayRef<uint8_t> sectionData) {
1057 if (sectionData.empty())
1058 return success();
1059 EncodingReader propReader(sectionData, fileLoc);
1060 uint64_t count;
1061 if (failed(Result: propReader.parseVarInt(result&: count)))
1062 return failure();
1063 // Parse the raw properties buffer.
1064 if (failed(Result: propReader.parseBytes(length: propReader.size(), result&: propertiesBuffers)))
1065 return failure();
1066
1067 EncodingReader offsetsReader(propertiesBuffers, fileLoc);
1068 offsetTable.reserve(N: count);
1069 for (auto idx : llvm::seq<int64_t>(Begin: 0, End: count)) {
1070 (void)idx;
1071 offsetTable.push_back(Elt: propertiesBuffers.size() - offsetsReader.size());
1072 ArrayRef<uint8_t> rawProperties;
1073 uint64_t dataSize;
1074 if (failed(Result: offsetsReader.parseVarInt(result&: dataSize)) ||
1075 failed(Result: offsetsReader.parseBytes(length: dataSize, result&: rawProperties)))
1076 return failure();
1077 }
1078 if (!offsetsReader.empty())
1079 return offsetsReader.emitError()
1080 << "Broken properties section: didn't exhaust the offsets table";
1081 return success();
1082 }
1083
1084 LogicalResult read(Location fileLoc, DialectReader &dialectReader,
1085 OperationName *opName, OperationState &opState) const {
1086 uint64_t propertiesIdx;
1087 if (failed(Result: dialectReader.readVarInt(result&: propertiesIdx)))
1088 return failure();
1089 if (propertiesIdx >= offsetTable.size())
1090 return dialectReader.emitError(msg: "Properties idx out-of-bound for ")
1091 << opName->getStringRef();
1092 size_t propertiesOffset = offsetTable[propertiesIdx];
1093 if (propertiesIdx >= propertiesBuffers.size())
1094 return dialectReader.emitError(msg: "Properties offset out-of-bound for ")
1095 << opName->getStringRef();
1096
1097 // Acquire the sub-buffer that represent the requested properties.
1098 ArrayRef<char> rawProperties;
1099 {
1100 // "Seek" to the requested offset by getting a new reader with the right
1101 // sub-buffer.
1102 EncodingReader reader(propertiesBuffers.drop_front(N: propertiesOffset),
1103 fileLoc);
1104 // Properties are stored as a sequence of {size + raw_data}.
1105 if (failed(
1106 Result: dialectReader.withEncodingReader(encReader&: reader).readBlob(result&: rawProperties)))
1107 return failure();
1108 }
1109 // Setup a new reader to read from the `rawProperties` sub-buffer.
1110 EncodingReader reader(
1111 StringRef(rawProperties.begin(), rawProperties.size()), fileLoc);
1112 DialectReader propReader = dialectReader.withEncodingReader(encReader&: reader);
1113
1114 auto *iface = opName->getInterface<BytecodeOpInterface>();
1115 if (iface)
1116 return iface->readProperties(propReader, opState);
1117 if (opName->isRegistered())
1118 return propReader.emitError(
1119 msg: "has properties but missing BytecodeOpInterface for ")
1120 << opName->getStringRef();
1121 // Unregistered op are storing properties as an attribute.
1122 return propReader.readAttribute(result&: opState.propertiesAttr);
1123 }
1124
1125private:
1126 /// The properties buffer referenced within the bytecode file.
1127 ArrayRef<uint8_t> propertiesBuffers;
1128
1129 /// Table of offset in the buffer above.
1130 SmallVector<int64_t> offsetTable;
1131};
1132} // namespace
1133
1134LogicalResult AttrTypeReader::initialize(
1135 MutableArrayRef<std::unique_ptr<BytecodeDialect>> dialects,
1136 ArrayRef<uint8_t> sectionData, ArrayRef<uint8_t> offsetSectionData) {
1137 EncodingReader offsetReader(offsetSectionData, fileLoc);
1138
1139 // Parse the number of attribute and type entries.
1140 uint64_t numAttributes, numTypes;
1141 if (failed(Result: offsetReader.parseVarInt(result&: numAttributes)) ||
1142 failed(Result: offsetReader.parseVarInt(result&: numTypes)))
1143 return failure();
1144 attributes.resize(N: numAttributes);
1145 types.resize(N: numTypes);
1146
1147 // A functor used to accumulate the offsets for the entries in the given
1148 // range.
1149 uint64_t currentOffset = 0;
1150 auto parseEntries = [&](auto &&range) {
1151 size_t currentIndex = 0, endIndex = range.size();
1152
1153 // Parse an individual entry.
1154 auto parseEntryFn = [&](BytecodeDialect *dialect) -> LogicalResult {
1155 auto &entry = range[currentIndex++];
1156
1157 uint64_t entrySize;
1158 if (failed(offsetReader.parseVarIntWithFlag(result&: entrySize,
1159 flag&: entry.hasCustomEncoding)))
1160 return failure();
1161
1162 // Verify that the offset is actually valid.
1163 if (currentOffset + entrySize > sectionData.size()) {
1164 return offsetReader.emitError(
1165 args: "Attribute or Type entry offset points past the end of section");
1166 }
1167
1168 entry.data = sectionData.slice(N: currentOffset, M: entrySize);
1169 entry.dialect = dialect;
1170 currentOffset += entrySize;
1171 return success();
1172 };
1173 while (currentIndex != endIndex)
1174 if (failed(parseDialectGrouping(offsetReader, dialects, parseEntryFn)))
1175 return failure();
1176 return success();
1177 };
1178
1179 // Process each of the attributes, and then the types.
1180 if (failed(Result: parseEntries(attributes)) || failed(Result: parseEntries(types)))
1181 return failure();
1182
1183 // Ensure that we read everything from the section.
1184 if (!offsetReader.empty()) {
1185 return offsetReader.emitError(
1186 args: "unexpected trailing data in the Attribute/Type offset section");
1187 }
1188
1189 return success();
1190}
1191
1192template <typename T>
1193T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
1194 StringRef entryType) {
1195 if (index >= entries.size()) {
1196 emitError(loc: fileLoc) << "invalid " << entryType << " index: " << index;
1197 return {};
1198 }
1199
1200 // If the entry has already been resolved, there is nothing left to do.
1201 Entry<T> &entry = entries[index];
1202 if (entry.entry)
1203 return entry.entry;
1204
1205 // Parse the entry.
1206 EncodingReader reader(entry.data, fileLoc);
1207
1208 // Parse based on how the entry was encoded.
1209 if (entry.hasCustomEncoding) {
1210 if (failed(parseCustomEntry(entry, reader, entryType)))
1211 return T();
1212 } else if (failed(parseAsmEntry(entry.entry, reader, entryType))) {
1213 return T();
1214 }
1215
1216 if (!reader.empty()) {
1217 reader.emitError(args: "unexpected trailing bytes after " + entryType + " entry");
1218 return T();
1219 }
1220 return entry.entry;
1221}
1222
1223template <typename T>
1224LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
1225 StringRef entryType) {
1226 StringRef asmStr;
1227 if (failed(Result: reader.parseNullTerminatedString(result&: asmStr)))
1228 return failure();
1229
1230 // Invoke the MLIR assembly parser to parse the entry text.
1231 size_t numRead = 0;
1232 MLIRContext *context = fileLoc->getContext();
1233 if constexpr (std::is_same_v<T, Type>)
1234 result =
1235 ::parseType(typeStr: asmStr, context, numRead: &numRead, /*isKnownNullTerminated=*/true);
1236 else
1237 result = ::parseAttribute(attrStr: asmStr, context, type: Type(), numRead: &numRead,
1238 /*isKnownNullTerminated=*/true);
1239 if (!result)
1240 return failure();
1241
1242 // Ensure there weren't dangling characters after the entry.
1243 if (numRead != asmStr.size()) {
1244 return reader.emitError(args: "trailing characters found after ", args&: entryType,
1245 args: " assembly format: ", args: asmStr.drop_front(N: numRead));
1246 }
1247 return success();
1248}
1249
1250template <typename T>
1251LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
1252 EncodingReader &reader,
1253 StringRef entryType) {
1254 DialectReader dialectReader(*this, stringReader, resourceReader, dialectsMap,
1255 reader, bytecodeVersion);
1256 if (failed(entry.dialect->load(dialectReader, fileLoc.getContext())))
1257 return failure();
1258
1259 if constexpr (std::is_same_v<T, Type>) {
1260 // Try parsing with callbacks first if available.
1261 for (const auto &callback :
1262 parserConfig.getBytecodeReaderConfig().getTypeCallbacks()) {
1263 if (failed(
1264 callback->read(reader&: dialectReader, dialectName: entry.dialect->name, entry&: entry.entry)))
1265 return failure();
1266 // Early return if parsing was successful.
1267 if (!!entry.entry)
1268 return success();
1269
1270 // Reset the reader if we failed to parse, so we can fall through the
1271 // other parsing functions.
1272 reader = EncodingReader(entry.data, reader.getLoc());
1273 }
1274 } else {
1275 // Try parsing with callbacks first if available.
1276 for (const auto &callback :
1277 parserConfig.getBytecodeReaderConfig().getAttributeCallbacks()) {
1278 if (failed(
1279 callback->read(reader&: dialectReader, dialectName: entry.dialect->name, entry&: entry.entry)))
1280 return failure();
1281 // Early return if parsing was successful.
1282 if (!!entry.entry)
1283 return success();
1284
1285 // Reset the reader if we failed to parse, so we can fall through the
1286 // other parsing functions.
1287 reader = EncodingReader(entry.data, reader.getLoc());
1288 }
1289 }
1290
1291 // Ensure that the dialect implements the bytecode interface.
1292 if (!entry.dialect->interface) {
1293 return reader.emitError("dialect '", entry.dialect->name,
1294 "' does not implement the bytecode interface");
1295 }
1296
1297 if constexpr (std::is_same_v<T, Type>)
1298 entry.entry = entry.dialect->interface->readType(dialectReader);
1299 else
1300 entry.entry = entry.dialect->interface->readAttribute(dialectReader);
1301
1302 return success(!!entry.entry);
1303}
1304
1305//===----------------------------------------------------------------------===//
1306// Bytecode Reader
1307//===----------------------------------------------------------------------===//
1308
1309/// This class is used to read a bytecode buffer and translate it into MLIR.
1310class mlir::BytecodeReader::Impl {
1311 struct RegionReadState;
1312 using LazyLoadableOpsInfo =
1313 std::list<std::pair<Operation *, RegionReadState>>;
1314 using LazyLoadableOpsMap =
1315 DenseMap<Operation *, LazyLoadableOpsInfo::iterator>;
1316
1317public:
1318 Impl(Location fileLoc, const ParserConfig &config, bool lazyLoading,
1319 llvm::MemoryBufferRef buffer,
1320 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef)
1321 : config(config), fileLoc(fileLoc), lazyLoading(lazyLoading),
1322 attrTypeReader(stringReader, resourceReader, dialectsMap, version,
1323 fileLoc, config),
1324 // Use the builtin unrealized conversion cast operation to represent
1325 // forward references to values that aren't yet defined.
1326 forwardRefOpState(UnknownLoc::get(config.getContext()),
1327 "builtin.unrealized_conversion_cast", ValueRange(),
1328 NoneType::get(config.getContext())),
1329 buffer(buffer), bufferOwnerRef(bufferOwnerRef) {}
1330
1331 /// Read the bytecode defined within `buffer` into the given block.
1332 LogicalResult read(Block *block,
1333 llvm::function_ref<bool(Operation *)> lazyOps);
1334
1335 /// Return the number of ops that haven't been materialized yet.
1336 int64_t getNumOpsToMaterialize() const { return lazyLoadableOpsMap.size(); }
1337
1338 bool isMaterializable(Operation *op) { return lazyLoadableOpsMap.count(Val: op); }
1339
1340 /// Materialize the provided operation, invoke the lazyOpsCallback on every
1341 /// newly found lazy operation.
1342 LogicalResult
1343 materialize(Operation *op,
1344 llvm::function_ref<bool(Operation *)> lazyOpsCallback) {
1345 this->lazyOpsCallback = lazyOpsCallback;
1346 auto resetlazyOpsCallback =
1347 llvm::make_scope_exit(F: [&] { this->lazyOpsCallback = nullptr; });
1348 auto it = lazyLoadableOpsMap.find(Val: op);
1349 assert(it != lazyLoadableOpsMap.end() &&
1350 "materialize called on non-materializable op");
1351 return materialize(it);
1352 }
1353
1354 /// Materialize all operations.
1355 LogicalResult materializeAll() {
1356 while (!lazyLoadableOpsMap.empty()) {
1357 if (failed(Result: materialize(it: lazyLoadableOpsMap.begin())))
1358 return failure();
1359 }
1360 return success();
1361 }
1362
1363 /// Finalize the lazy-loading by calling back with every op that hasn't been
1364 /// materialized to let the client decide if the op should be deleted or
1365 /// materialized. The op is materialized if the callback returns true, deleted
1366 /// otherwise.
1367 LogicalResult finalize(function_ref<bool(Operation *)> shouldMaterialize) {
1368 while (!lazyLoadableOps.empty()) {
1369 Operation *op = lazyLoadableOps.begin()->first;
1370 if (shouldMaterialize(op)) {
1371 if (failed(Result: materialize(it: lazyLoadableOpsMap.find(Val: op))))
1372 return failure();
1373 continue;
1374 }
1375 op->dropAllReferences();
1376 op->erase();
1377 lazyLoadableOps.pop_front();
1378 lazyLoadableOpsMap.erase(Val: op);
1379 }
1380 return success();
1381 }
1382
1383private:
1384 LogicalResult materialize(LazyLoadableOpsMap::iterator it) {
1385 assert(it != lazyLoadableOpsMap.end() &&
1386 "materialize called on non-materializable op");
1387 valueScopes.emplace_back();
1388 std::vector<RegionReadState> regionStack;
1389 regionStack.push_back(x: std::move(it->getSecond()->second));
1390 lazyLoadableOps.erase(position: it->getSecond());
1391 lazyLoadableOpsMap.erase(I: it);
1392
1393 while (!regionStack.empty())
1394 if (failed(Result: parseRegions(regionStack, readState&: regionStack.back())))
1395 return failure();
1396 return success();
1397 }
1398
1399 /// Return the context for this config.
1400 MLIRContext *getContext() const { return config.getContext(); }
1401
1402 /// Parse the bytecode version.
1403 LogicalResult parseVersion(EncodingReader &reader);
1404
1405 //===--------------------------------------------------------------------===//
1406 // Dialect Section
1407
1408 LogicalResult parseDialectSection(ArrayRef<uint8_t> sectionData);
1409
1410 /// Parse an operation name reference using the given reader, and set the
1411 /// `wasRegistered` flag that indicates if the bytecode was produced by a
1412 /// context where opName was registered.
1413 FailureOr<OperationName> parseOpName(EncodingReader &reader,
1414 std::optional<bool> &wasRegistered);
1415
1416 //===--------------------------------------------------------------------===//
1417 // Attribute/Type Section
1418
1419 /// Parse an attribute or type using the given reader.
1420 template <typename T>
1421 LogicalResult parseAttribute(EncodingReader &reader, T &result) {
1422 return attrTypeReader.parseAttribute(reader, result);
1423 }
1424 LogicalResult parseType(EncodingReader &reader, Type &result) {
1425 return attrTypeReader.parseType(reader, result);
1426 }
1427
1428 //===--------------------------------------------------------------------===//
1429 // Resource Section
1430
1431 LogicalResult
1432 parseResourceSection(EncodingReader &reader,
1433 std::optional<ArrayRef<uint8_t>> resourceData,
1434 std::optional<ArrayRef<uint8_t>> resourceOffsetData);
1435
1436 //===--------------------------------------------------------------------===//
1437 // IR Section
1438
1439 /// This struct represents the current read state of a range of regions. This
1440 /// struct is used to enable iterative parsing of regions.
1441 struct RegionReadState {
1442 RegionReadState(Operation *op, EncodingReader *reader,
1443 bool isIsolatedFromAbove)
1444 : RegionReadState(op->getRegions(), reader, isIsolatedFromAbove) {}
1445 RegionReadState(MutableArrayRef<Region> regions, EncodingReader *reader,
1446 bool isIsolatedFromAbove)
1447 : curRegion(regions.begin()), endRegion(regions.end()), reader(reader),
1448 isIsolatedFromAbove(isIsolatedFromAbove) {}
1449
1450 /// The current regions being read.
1451 MutableArrayRef<Region>::iterator curRegion, endRegion;
1452 /// This is the reader to use for this region, this pointer is pointing to
1453 /// the parent region reader unless the current region is IsolatedFromAbove,
1454 /// in which case the pointer is pointing to the `owningReader` which is a
1455 /// section dedicated to the current region.
1456 EncodingReader *reader;
1457 std::unique_ptr<EncodingReader> owningReader;
1458
1459 /// The number of values defined immediately within this region.
1460 unsigned numValues = 0;
1461
1462 /// The current blocks of the region being read.
1463 SmallVector<Block *> curBlocks;
1464 Region::iterator curBlock = {};
1465
1466 /// The number of operations remaining to be read from the current block
1467 /// being read.
1468 uint64_t numOpsRemaining = 0;
1469
1470 /// A flag indicating if the regions being read are isolated from above.
1471 bool isIsolatedFromAbove = false;
1472 };
1473
1474 LogicalResult parseIRSection(ArrayRef<uint8_t> sectionData, Block *block);
1475 LogicalResult parseRegions(std::vector<RegionReadState> &regionStack,
1476 RegionReadState &readState);
1477 FailureOr<Operation *> parseOpWithoutRegions(EncodingReader &reader,
1478 RegionReadState &readState,
1479 bool &isIsolatedFromAbove);
1480
1481 LogicalResult parseRegion(RegionReadState &readState);
1482 LogicalResult parseBlockHeader(EncodingReader &reader,
1483 RegionReadState &readState);
1484 LogicalResult parseBlockArguments(EncodingReader &reader, Block *block);
1485
1486 //===--------------------------------------------------------------------===//
1487 // Value Processing
1488
1489 /// Parse an operand reference using the given reader. Returns nullptr in the
1490 /// case of failure.
1491 Value parseOperand(EncodingReader &reader);
1492
1493 /// Sequentially define the given value range.
1494 LogicalResult defineValues(EncodingReader &reader, ValueRange values);
1495
1496 /// Create a value to use for a forward reference.
1497 Value createForwardRef();
1498
1499 //===--------------------------------------------------------------------===//
1500 // Use-list order helpers
1501
1502 /// This struct is a simple storage that contains information required to
1503 /// reorder the use-list of a value with respect to the pre-order traversal
1504 /// ordering.
1505 struct UseListOrderStorage {
1506 UseListOrderStorage(bool isIndexPairEncoding,
1507 SmallVector<unsigned, 4> &&indices)
1508 : indices(std::move(indices)),
1509 isIndexPairEncoding(isIndexPairEncoding){};
1510 /// The vector containing the information required to reorder the
1511 /// use-list of a value.
1512 SmallVector<unsigned, 4> indices;
1513
1514 /// Whether indices represent a pair of type `(src, dst)` or it is a direct
1515 /// indexing, such as `dst = order[src]`.
1516 bool isIndexPairEncoding;
1517 };
1518
1519 /// Parse use-list order from bytecode for a range of values if available. The
1520 /// range is expected to be either a block argument or an op result range. On
1521 /// success, return a map of the position in the range and the use-list order
1522 /// encoding. The function assumes to know the size of the range it is
1523 /// processing.
1524 using UseListMapT = DenseMap<unsigned, UseListOrderStorage>;
1525 FailureOr<UseListMapT> parseUseListOrderForRange(EncodingReader &reader,
1526 uint64_t rangeSize);
1527
1528 /// Shuffle the use-chain according to the order parsed.
1529 LogicalResult sortUseListOrder(Value value);
1530
1531 /// Recursively visit all the values defined within topLevelOp and sort the
1532 /// use-list orders according to the indices parsed.
1533 LogicalResult processUseLists(Operation *topLevelOp);
1534
1535 //===--------------------------------------------------------------------===//
1536 // Fields
1537
1538 /// This class represents a single value scope, in which a value scope is
1539 /// delimited by isolated from above regions.
1540 struct ValueScope {
1541 /// Push a new region state onto this scope, reserving enough values for
1542 /// those defined within the current region of the provided state.
1543 void push(RegionReadState &readState) {
1544 nextValueIDs.push_back(Elt: values.size());
1545 values.resize(new_size: values.size() + readState.numValues);
1546 }
1547
1548 /// Pop the values defined for the current region within the provided region
1549 /// state.
1550 void pop(RegionReadState &readState) {
1551 values.resize(new_size: values.size() - readState.numValues);
1552 nextValueIDs.pop_back();
1553 }
1554
1555 /// The set of values defined in this scope.
1556 std::vector<Value> values;
1557
1558 /// The ID for the next defined value for each region current being
1559 /// processed in this scope.
1560 SmallVector<unsigned, 4> nextValueIDs;
1561 };
1562
1563 /// The configuration of the parser.
1564 const ParserConfig &config;
1565
1566 /// A location to use when emitting errors.
1567 Location fileLoc;
1568
1569 /// Flag that indicates if lazyloading is enabled.
1570 bool lazyLoading;
1571
1572 /// Keep track of operations that have been lazy loaded (their regions haven't
1573 /// been materialized), along with the `RegionReadState` that allows to
1574 /// lazy-load the regions nested under the operation.
1575 LazyLoadableOpsInfo lazyLoadableOps;
1576 LazyLoadableOpsMap lazyLoadableOpsMap;
1577 llvm::function_ref<bool(Operation *)> lazyOpsCallback;
1578
1579 /// The reader used to process attribute and types within the bytecode.
1580 AttrTypeReader attrTypeReader;
1581
1582 /// The version of the bytecode being read.
1583 uint64_t version = 0;
1584
1585 /// The producer of the bytecode being read.
1586 StringRef producer;
1587
1588 /// The table of IR units referenced within the bytecode file.
1589 SmallVector<std::unique_ptr<BytecodeDialect>> dialects;
1590 llvm::StringMap<BytecodeDialect *> dialectsMap;
1591 SmallVector<BytecodeOperationName> opNames;
1592
1593 /// The reader used to process resources within the bytecode.
1594 ResourceSectionReader resourceReader;
1595
1596 /// Worklist of values with custom use-list orders to process before the end
1597 /// of the parsing.
1598 DenseMap<void *, UseListOrderStorage> valueToUseListMap;
1599
1600 /// The table of strings referenced within the bytecode file.
1601 StringSectionReader stringReader;
1602
1603 /// The table of properties referenced by the operation in the bytecode file.
1604 PropertiesSectionReader propertiesReader;
1605
1606 /// The current set of available IR value scopes.
1607 std::vector<ValueScope> valueScopes;
1608
1609 /// The global pre-order operation ordering.
1610 DenseMap<Operation *, unsigned> operationIDs;
1611
1612 /// A block containing the set of operations defined to create forward
1613 /// references.
1614 Block forwardRefOps;
1615
1616 /// A block containing previously created, and no longer used, forward
1617 /// reference operations.
1618 Block openForwardRefOps;
1619
1620 /// An operation state used when instantiating forward references.
1621 OperationState forwardRefOpState;
1622
1623 /// Reference to the input buffer.
1624 llvm::MemoryBufferRef buffer;
1625
1626 /// The optional owning source manager, which when present may be used to
1627 /// extend the lifetime of the input buffer.
1628 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef;
1629};
1630
1631LogicalResult BytecodeReader::Impl::read(
1632 Block *block, llvm::function_ref<bool(Operation *)> lazyOpsCallback) {
1633 EncodingReader reader(buffer.getBuffer(), fileLoc);
1634 this->lazyOpsCallback = lazyOpsCallback;
1635 auto resetlazyOpsCallback =
1636 llvm::make_scope_exit(F: [&] { this->lazyOpsCallback = nullptr; });
1637
1638 // Skip over the bytecode header, this should have already been checked.
1639 if (failed(Result: reader.skipBytes(length: StringRef("ML\xefR").size())))
1640 return failure();
1641 // Parse the bytecode version and producer.
1642 if (failed(Result: parseVersion(reader)) ||
1643 failed(Result: reader.parseNullTerminatedString(result&: producer)))
1644 return failure();
1645
1646 // Add a diagnostic handler that attaches a note that includes the original
1647 // producer of the bytecode.
1648 ScopedDiagnosticHandler diagHandler(getContext(), [&](Diagnostic &diag) {
1649 diag.attachNote() << "in bytecode version " << version
1650 << " produced by: " << producer;
1651 return failure();
1652 });
1653
1654 // Parse the raw data for each of the top-level sections of the bytecode.
1655 std::optional<ArrayRef<uint8_t>>
1656 sectionDatas[bytecode::Section::kNumSections];
1657 while (!reader.empty()) {
1658 // Read the next section from the bytecode.
1659 bytecode::Section::ID sectionID;
1660 ArrayRef<uint8_t> sectionData;
1661 if (failed(Result: reader.parseSection(sectionID, sectionData)))
1662 return failure();
1663
1664 // Check for duplicate sections, we only expect one instance of each.
1665 if (sectionDatas[sectionID]) {
1666 return reader.emitError(args: "duplicate top-level section: ",
1667 args: ::toString(sectionID));
1668 }
1669 sectionDatas[sectionID] = sectionData;
1670 }
1671 // Check that all of the required sections were found.
1672 for (int i = 0; i < bytecode::Section::kNumSections; ++i) {
1673 bytecode::Section::ID sectionID = static_cast<bytecode::Section::ID>(i);
1674 if (!sectionDatas[i] && !isSectionOptional(sectionID, version)) {
1675 return reader.emitError(args: "missing data for top-level section: ",
1676 args: ::toString(sectionID));
1677 }
1678 }
1679
1680 // Process the string section first.
1681 if (failed(Result: stringReader.initialize(
1682 fileLoc, sectionData: *sectionDatas[bytecode::Section::kString])))
1683 return failure();
1684
1685 // Process the properties section.
1686 if (sectionDatas[bytecode::Section::kProperties] &&
1687 failed(Result: propertiesReader.initialize(
1688 fileLoc, sectionData: *sectionDatas[bytecode::Section::kProperties])))
1689 return failure();
1690
1691 // Process the dialect section.
1692 if (failed(Result: parseDialectSection(sectionData: *sectionDatas[bytecode::Section::kDialect])))
1693 return failure();
1694
1695 // Process the resource section if present.
1696 if (failed(Result: parseResourceSection(
1697 reader, resourceData: sectionDatas[bytecode::Section::kResource],
1698 resourceOffsetData: sectionDatas[bytecode::Section::kResourceOffset])))
1699 return failure();
1700
1701 // Process the attribute and type section.
1702 if (failed(Result: attrTypeReader.initialize(
1703 dialects, sectionData: *sectionDatas[bytecode::Section::kAttrType],
1704 offsetSectionData: *sectionDatas[bytecode::Section::kAttrTypeOffset])))
1705 return failure();
1706
1707 // Finally, process the IR section.
1708 return parseIRSection(sectionData: *sectionDatas[bytecode::Section::kIR], block);
1709}
1710
1711LogicalResult BytecodeReader::Impl::parseVersion(EncodingReader &reader) {
1712 if (failed(Result: reader.parseVarInt(result&: version)))
1713 return failure();
1714
1715 // Validate the bytecode version.
1716 uint64_t currentVersion = bytecode::kVersion;
1717 uint64_t minSupportedVersion = bytecode::kMinSupportedVersion;
1718 if (version < minSupportedVersion) {
1719 return reader.emitError(args: "bytecode version ", args&: version,
1720 args: " is older than the current version of ",
1721 args&: currentVersion, args: ", and upgrade is not supported");
1722 }
1723 if (version > currentVersion) {
1724 return reader.emitError(args: "bytecode version ", args&: version,
1725 args: " is newer than the current version ",
1726 args&: currentVersion);
1727 }
1728 // Override any request to lazy-load if the bytecode version is too old.
1729 if (version < bytecode::kLazyLoading)
1730 lazyLoading = false;
1731 return success();
1732}
1733
1734//===----------------------------------------------------------------------===//
1735// Dialect Section
1736//===----------------------------------------------------------------------===//
1737
1738LogicalResult BytecodeDialect::load(const DialectReader &reader,
1739 MLIRContext *ctx) {
1740 if (dialect)
1741 return success();
1742 Dialect *loadedDialect = ctx->getOrLoadDialect(name);
1743 if (!loadedDialect && !ctx->allowsUnregisteredDialects()) {
1744 return reader.emitError(msg: "dialect '")
1745 << name
1746 << "' is unknown. If this is intended, please call "
1747 "allowUnregisteredDialects() on the MLIRContext, or use "
1748 "-allow-unregistered-dialect with the MLIR tool used.";
1749 }
1750 dialect = loadedDialect;
1751
1752 // If the dialect was actually loaded, check to see if it has a bytecode
1753 // interface.
1754 if (loadedDialect)
1755 interface = dyn_cast<BytecodeDialectInterface>(Val: loadedDialect);
1756 if (!versionBuffer.empty()) {
1757 if (!interface)
1758 return reader.emitError(msg: "dialect '")
1759 << name
1760 << "' does not implement the bytecode interface, "
1761 "but found a version entry";
1762 EncodingReader encReader(versionBuffer, reader.getLoc());
1763 DialectReader versionReader = reader.withEncodingReader(encReader);
1764 loadedVersion = interface->readVersion(reader&: versionReader);
1765 if (!loadedVersion)
1766 return failure();
1767 }
1768 return success();
1769}
1770
1771LogicalResult
1772BytecodeReader::Impl::parseDialectSection(ArrayRef<uint8_t> sectionData) {
1773 EncodingReader sectionReader(sectionData, fileLoc);
1774
1775 // Parse the number of dialects in the section.
1776 uint64_t numDialects;
1777 if (failed(Result: sectionReader.parseVarInt(result&: numDialects)))
1778 return failure();
1779 dialects.resize(N: numDialects);
1780
1781 // Parse each of the dialects.
1782 for (uint64_t i = 0; i < numDialects; ++i) {
1783 dialects[i] = std::make_unique<BytecodeDialect>();
1784 /// Before version kDialectVersioning, there wasn't any versioning available
1785 /// for dialects, and the entryIdx represent the string itself.
1786 if (version < bytecode::kDialectVersioning) {
1787 if (failed(Result: stringReader.parseString(reader&: sectionReader, result&: dialects[i]->name)))
1788 return failure();
1789 continue;
1790 }
1791
1792 // Parse ID representing dialect and version.
1793 uint64_t dialectNameIdx;
1794 bool versionAvailable;
1795 if (failed(Result: sectionReader.parseVarIntWithFlag(result&: dialectNameIdx,
1796 flag&: versionAvailable)))
1797 return failure();
1798 if (failed(Result: stringReader.parseStringAtIndex(reader&: sectionReader, index: dialectNameIdx,
1799 result&: dialects[i]->name)))
1800 return failure();
1801 if (versionAvailable) {
1802 bytecode::Section::ID sectionID;
1803 if (failed(Result: sectionReader.parseSection(sectionID,
1804 sectionData&: dialects[i]->versionBuffer)))
1805 return failure();
1806 if (sectionID != bytecode::Section::kDialectVersions) {
1807 emitError(loc: fileLoc, message: "expected dialect version section");
1808 return failure();
1809 }
1810 }
1811 dialectsMap[dialects[i]->name] = dialects[i].get();
1812 }
1813
1814 // Parse the operation names, which are grouped by dialect.
1815 auto parseOpName = [&](BytecodeDialect *dialect) {
1816 StringRef opName;
1817 std::optional<bool> wasRegistered;
1818 // Prior to version kNativePropertiesEncoding, the information about wheter
1819 // an op was registered or not wasn't encoded.
1820 if (version < bytecode::kNativePropertiesEncoding) {
1821 if (failed(Result: stringReader.parseString(reader&: sectionReader, result&: opName)))
1822 return failure();
1823 } else {
1824 bool wasRegisteredFlag;
1825 if (failed(Result: stringReader.parseStringWithFlag(reader&: sectionReader, result&: opName,
1826 flag&: wasRegisteredFlag)))
1827 return failure();
1828 wasRegistered = wasRegisteredFlag;
1829 }
1830 opNames.emplace_back(Args&: dialect, Args&: opName, Args&: wasRegistered);
1831 return success();
1832 };
1833 // Avoid re-allocation in bytecode version >=kElideUnknownBlockArgLocation
1834 // where the number of ops are known.
1835 if (version >= bytecode::kElideUnknownBlockArgLocation) {
1836 uint64_t numOps;
1837 if (failed(Result: sectionReader.parseVarInt(result&: numOps)))
1838 return failure();
1839 opNames.reserve(N: numOps);
1840 }
1841 while (!sectionReader.empty())
1842 if (failed(Result: parseDialectGrouping(reader&: sectionReader, dialects, entryCallback: parseOpName)))
1843 return failure();
1844 return success();
1845}
1846
1847FailureOr<OperationName>
1848BytecodeReader::Impl::parseOpName(EncodingReader &reader,
1849 std::optional<bool> &wasRegistered) {
1850 BytecodeOperationName *opName = nullptr;
1851 if (failed(Result: parseEntry(reader, entries&: opNames, entry&: opName, entryStr: "operation name")))
1852 return failure();
1853 wasRegistered = opName->wasRegistered;
1854 // Check to see if this operation name has already been resolved. If we
1855 // haven't, load the dialect and build the operation name.
1856 if (!opName->opName) {
1857 // If the opName is empty, this is because we use to accept names such as
1858 // `foo` without any `.` separator. We shouldn't tolerate this in textual
1859 // format anymore but for now we'll be backward compatible. This can only
1860 // happen with unregistered dialects.
1861 if (opName->name.empty()) {
1862 opName->opName.emplace(args&: opName->dialect->name, args: getContext());
1863 } else {
1864 // Load the dialect and its version.
1865 DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
1866 dialectsMap, reader, version);
1867 if (failed(Result: opName->dialect->load(reader: dialectReader, ctx: getContext())))
1868 return failure();
1869 opName->opName.emplace(args: (opName->dialect->name + "." + opName->name).str(),
1870 args: getContext());
1871 }
1872 }
1873 return *opName->opName;
1874}
1875
1876//===----------------------------------------------------------------------===//
1877// Resource Section
1878//===----------------------------------------------------------------------===//
1879
1880LogicalResult BytecodeReader::Impl::parseResourceSection(
1881 EncodingReader &reader, std::optional<ArrayRef<uint8_t>> resourceData,
1882 std::optional<ArrayRef<uint8_t>> resourceOffsetData) {
1883 // Ensure both sections are either present or not.
1884 if (resourceData.has_value() != resourceOffsetData.has_value()) {
1885 if (resourceOffsetData)
1886 return emitError(loc: fileLoc, message: "unexpected resource offset section when "
1887 "resource section is not present");
1888 return emitError(
1889 loc: fileLoc,
1890 message: "expected resource offset section when resource section is present");
1891 }
1892
1893 // If the resource sections are absent, there is nothing to do.
1894 if (!resourceData)
1895 return success();
1896
1897 // Initialize the resource reader with the resource sections.
1898 DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
1899 dialectsMap, reader, version);
1900 return resourceReader.initialize(fileLoc, config, dialects, stringReader,
1901 sectionData: *resourceData, offsetSectionData: *resourceOffsetData,
1902 dialectReader, bufferOwnerRef);
1903}
1904
1905//===----------------------------------------------------------------------===//
1906// UseListOrder Helpers
1907//===----------------------------------------------------------------------===//
1908
1909FailureOr<BytecodeReader::Impl::UseListMapT>
1910BytecodeReader::Impl::parseUseListOrderForRange(EncodingReader &reader,
1911 uint64_t numResults) {
1912 BytecodeReader::Impl::UseListMapT map;
1913 uint64_t numValuesToRead = 1;
1914 if (numResults > 1 && failed(Result: reader.parseVarInt(result&: numValuesToRead)))
1915 return failure();
1916
1917 for (size_t valueIdx = 0; valueIdx < numValuesToRead; valueIdx++) {
1918 uint64_t resultIdx = 0;
1919 if (numResults > 1 && failed(Result: reader.parseVarInt(result&: resultIdx)))
1920 return failure();
1921
1922 uint64_t numValues;
1923 bool indexPairEncoding;
1924 if (failed(Result: reader.parseVarIntWithFlag(result&: numValues, flag&: indexPairEncoding)))
1925 return failure();
1926
1927 SmallVector<unsigned, 4> useListOrders;
1928 for (size_t idx = 0; idx < numValues; idx++) {
1929 uint64_t index;
1930 if (failed(Result: reader.parseVarInt(result&: index)))
1931 return failure();
1932 useListOrders.push_back(Elt: index);
1933 }
1934
1935 // Store in a map the result index
1936 map.try_emplace(Key: resultIdx, Args: UseListOrderStorage(indexPairEncoding,
1937 std::move(useListOrders)));
1938 }
1939
1940 return map;
1941}
1942
1943/// Sorts each use according to the order specified in the use-list parsed. If
1944/// the custom use-list is not found, this means that the order needs to be
1945/// consistent with the reverse pre-order walk of the IR. If multiple uses lie
1946/// on the same operation, the order will follow the reverse operand number
1947/// ordering.
1948LogicalResult BytecodeReader::Impl::sortUseListOrder(Value value) {
1949 // Early return for trivial use-lists.
1950 if (value.use_empty() || value.hasOneUse())
1951 return success();
1952
1953 bool hasIncomingOrder =
1954 valueToUseListMap.contains(Val: value.getAsOpaquePointer());
1955
1956 // Compute the current order of the use-list with respect to the global
1957 // ordering. Detect if the order is already sorted while doing so.
1958 bool alreadySorted = true;
1959 auto &firstUse = *value.use_begin();
1960 uint64_t prevID =
1961 bytecode::getUseID(val&: firstUse, ownerID: operationIDs.at(Val: firstUse.getOwner()));
1962 llvm::SmallVector<std::pair<unsigned, uint64_t>> currentOrder = {{0, prevID}};
1963 for (auto item : llvm::drop_begin(RangeOrContainer: llvm::enumerate(First: value.getUses()))) {
1964 uint64_t currentID = bytecode::getUseID(
1965 val&: item.value(), ownerID: operationIDs.at(Val: item.value().getOwner()));
1966 alreadySorted &= prevID > currentID;
1967 currentOrder.push_back(Elt: {item.index(), currentID});
1968 prevID = currentID;
1969 }
1970
1971 // If the order is already sorted, and there wasn't a custom order to apply
1972 // from the bytecode file, we are done.
1973 if (alreadySorted && !hasIncomingOrder)
1974 return success();
1975
1976 // If not already sorted, sort the indices of the current order by descending
1977 // useIDs.
1978 if (!alreadySorted)
1979 std::sort(
1980 first: currentOrder.begin(), last: currentOrder.end(),
1981 comp: [](auto elem1, auto elem2) { return elem1.second > elem2.second; });
1982
1983 if (!hasIncomingOrder) {
1984 // If the bytecode file did not contain any custom use-list order, it means
1985 // that the order was descending useID. Hence, shuffle by the first index
1986 // of the `currentOrder` pair.
1987 SmallVector<unsigned> shuffle(llvm::make_first_range(c&: currentOrder));
1988 value.shuffleUseList(indices: shuffle);
1989 return success();
1990 }
1991
1992 // Pull the custom order info from the map.
1993 UseListOrderStorage customOrder =
1994 valueToUseListMap.at(Val: value.getAsOpaquePointer());
1995 SmallVector<unsigned, 4> shuffle = std::move(customOrder.indices);
1996 uint64_t numUses = value.getNumUses();
1997
1998 // If the encoding was a pair of indices `(src, dst)` for every permutation,
1999 // reconstruct the shuffle vector for every use. Initialize the shuffle vector
2000 // as identity, and then apply the mapping encoded in the indices.
2001 if (customOrder.isIndexPairEncoding) {
2002 // Return failure if the number of indices was not representing pairs.
2003 if (shuffle.size() & 1)
2004 return failure();
2005
2006 SmallVector<unsigned, 4> newShuffle(numUses);
2007 size_t idx = 0;
2008 std::iota(first: newShuffle.begin(), last: newShuffle.end(), value: idx);
2009 for (idx = 0; idx < shuffle.size(); idx += 2)
2010 newShuffle[shuffle[idx]] = shuffle[idx + 1];
2011
2012 shuffle = std::move(newShuffle);
2013 }
2014
2015 // Make sure that the indices represent a valid mapping. That is, the sum of
2016 // all the values needs to be equal to (numUses - 1) * numUses / 2, and no
2017 // duplicates are allowed in the list.
2018 DenseSet<unsigned> set;
2019 uint64_t accumulator = 0;
2020 for (const auto &elem : shuffle) {
2021 if (!set.insert(V: elem).second)
2022 return failure();
2023 accumulator += elem;
2024 }
2025 if (numUses != shuffle.size() ||
2026 accumulator != (((numUses - 1) * numUses) >> 1))
2027 return failure();
2028
2029 // Apply the current ordering map onto the shuffle vector to get the final
2030 // use-list sorting indices before shuffling.
2031 shuffle = SmallVector<unsigned, 4>(llvm::map_range(
2032 C&: currentOrder, F: [&](auto item) { return shuffle[item.first]; }));
2033 value.shuffleUseList(indices: shuffle);
2034 return success();
2035}
2036
2037LogicalResult BytecodeReader::Impl::processUseLists(Operation *topLevelOp) {
2038 // Precompute operation IDs according to the pre-order walk of the IR. We
2039 // can't do this while parsing since parseRegions ordering is not strictly
2040 // equal to the pre-order walk.
2041 unsigned operationID = 0;
2042 topLevelOp->walk<mlir::WalkOrder::PreOrder>(
2043 callback: [&](Operation *op) { operationIDs.try_emplace(Key: op, Args: operationID++); });
2044
2045 auto blockWalk = topLevelOp->walk(callback: [this](Block *block) {
2046 for (auto arg : block->getArguments())
2047 if (failed(Result: sortUseListOrder(value: arg)))
2048 return WalkResult::interrupt();
2049 return WalkResult::advance();
2050 });
2051
2052 auto resultWalk = topLevelOp->walk(callback: [this](Operation *op) {
2053 for (auto result : op->getResults())
2054 if (failed(Result: sortUseListOrder(value: result)))
2055 return WalkResult::interrupt();
2056 return WalkResult::advance();
2057 });
2058
2059 return failure(IsFailure: blockWalk.wasInterrupted() || resultWalk.wasInterrupted());
2060}
2061
2062//===----------------------------------------------------------------------===//
2063// IR Section
2064//===----------------------------------------------------------------------===//
2065
2066LogicalResult
2067BytecodeReader::Impl::parseIRSection(ArrayRef<uint8_t> sectionData,
2068 Block *block) {
2069 EncodingReader reader(sectionData, fileLoc);
2070
2071 // A stack of operation regions currently being read from the bytecode.
2072 std::vector<RegionReadState> regionStack;
2073
2074 // Parse the top-level block using a temporary module operation.
2075 OwningOpRef<ModuleOp> moduleOp = ModuleOp::create(fileLoc);
2076 regionStack.emplace_back(*moduleOp, &reader, /*isIsolatedFromAbove=*/true);
2077 regionStack.back().curBlocks.push_back(Elt: moduleOp->getBody());
2078 regionStack.back().curBlock = regionStack.back().curRegion->begin();
2079 if (failed(Result: parseBlockHeader(reader, readState&: regionStack.back())))
2080 return failure();
2081 valueScopes.emplace_back();
2082 valueScopes.back().push(readState&: regionStack.back());
2083
2084 // Iteratively parse regions until everything has been resolved.
2085 while (!regionStack.empty())
2086 if (failed(Result: parseRegions(regionStack, readState&: regionStack.back())))
2087 return failure();
2088 if (!forwardRefOps.empty()) {
2089 return reader.emitError(
2090 args: "not all forward unresolved forward operand references");
2091 }
2092
2093 // Sort use-lists according to what specified in bytecode.
2094 if (failed(processUseLists(topLevelOp: *moduleOp)))
2095 return reader.emitError(
2096 args: "parsed use-list orders were invalid and could not be applied");
2097
2098 // Resolve dialect version.
2099 for (const std::unique_ptr<BytecodeDialect> &byteCodeDialect : dialects) {
2100 // Parsing is complete, give an opportunity to each dialect to visit the
2101 // IR and perform upgrades.
2102 if (!byteCodeDialect->loadedVersion)
2103 continue;
2104 if (byteCodeDialect->interface &&
2105 failed(byteCodeDialect->interface->upgradeFromVersion(
2106 topLevelOp: *moduleOp, version: *byteCodeDialect->loadedVersion)))
2107 return failure();
2108 }
2109
2110 // Verify that the parsed operations are valid.
2111 if (config.shouldVerifyAfterParse() && failed(verify(*moduleOp)))
2112 return failure();
2113
2114 // Splice the parsed operations over to the provided top-level block.
2115 auto &parsedOps = moduleOp->getBody()->getOperations();
2116 auto &destOps = block->getOperations();
2117 destOps.splice(destOps.end(), parsedOps, parsedOps.begin(), parsedOps.end());
2118 return success();
2119}
2120
2121LogicalResult
2122BytecodeReader::Impl::parseRegions(std::vector<RegionReadState> &regionStack,
2123 RegionReadState &readState) {
2124 // Process regions, blocks, and operations until the end or if a nested
2125 // region is encountered. In this case we push a new state in regionStack and
2126 // return, the processing of the current region will resume afterward.
2127 for (; readState.curRegion != readState.endRegion; ++readState.curRegion) {
2128 // If the current block hasn't been setup yet, parse the header for this
2129 // region. The current block is already setup when this function was
2130 // interrupted to recurse down in a nested region and we resume the current
2131 // block after processing the nested region.
2132 if (readState.curBlock == Region::iterator()) {
2133 if (failed(Result: parseRegion(readState)))
2134 return failure();
2135
2136 // If the region is empty, there is nothing to more to do.
2137 if (readState.curRegion->empty())
2138 continue;
2139 }
2140
2141 // Parse the blocks within the region.
2142 EncodingReader &reader = *readState.reader;
2143 do {
2144 while (readState.numOpsRemaining--) {
2145 // Read in the next operation. We don't read its regions directly, we
2146 // handle those afterwards as necessary.
2147 bool isIsolatedFromAbove = false;
2148 FailureOr<Operation *> op =
2149 parseOpWithoutRegions(reader, readState, isIsolatedFromAbove);
2150 if (failed(Result: op))
2151 return failure();
2152
2153 // If the op has regions, add it to the stack for processing and return:
2154 // we stop the processing of the current region and resume it after the
2155 // inner one is completed. Unless LazyLoading is activated in which case
2156 // nested region parsing is delayed.
2157 if ((*op)->getNumRegions()) {
2158 RegionReadState childState(*op, &reader, isIsolatedFromAbove);
2159
2160 // Isolated regions are encoded as a section in version 2 and above.
2161 if (version >= bytecode::kLazyLoading && isIsolatedFromAbove) {
2162 bytecode::Section::ID sectionID;
2163 ArrayRef<uint8_t> sectionData;
2164 if (failed(Result: reader.parseSection(sectionID, sectionData)))
2165 return failure();
2166 if (sectionID != bytecode::Section::kIR)
2167 return emitError(loc: fileLoc, message: "expected IR section for region");
2168 childState.owningReader =
2169 std::make_unique<EncodingReader>(args&: sectionData, args&: fileLoc);
2170 childState.reader = childState.owningReader.get();
2171
2172 // If the user has a callback set, they have the opportunity to
2173 // control lazyloading as we go.
2174 if (lazyLoading && (!lazyOpsCallback || !lazyOpsCallback(*op))) {
2175 lazyLoadableOps.emplace_back(args&: *op, args: std::move(childState));
2176 lazyLoadableOpsMap.try_emplace(Key: *op,
2177 Args: std::prev(x: lazyLoadableOps.end()));
2178 continue;
2179 }
2180 }
2181 regionStack.push_back(x: std::move(childState));
2182
2183 // If the op is isolated from above, push a new value scope.
2184 if (isIsolatedFromAbove)
2185 valueScopes.emplace_back();
2186 return success();
2187 }
2188 }
2189
2190 // Move to the next block of the region.
2191 if (++readState.curBlock == readState.curRegion->end())
2192 break;
2193 if (failed(Result: parseBlockHeader(reader, readState)))
2194 return failure();
2195 } while (true);
2196
2197 // Reset the current block and any values reserved for this region.
2198 readState.curBlock = {};
2199 valueScopes.back().pop(readState);
2200 }
2201
2202 // When the regions have been fully parsed, pop them off of the read stack. If
2203 // the regions were isolated from above, we also pop the last value scope.
2204 if (readState.isIsolatedFromAbove) {
2205 assert(!valueScopes.empty() && "Expect a valueScope after reading region");
2206 valueScopes.pop_back();
2207 }
2208 assert(!regionStack.empty() && "Expect a regionStack after reading region");
2209 regionStack.pop_back();
2210 return success();
2211}
2212
2213FailureOr<Operation *>
2214BytecodeReader::Impl::parseOpWithoutRegions(EncodingReader &reader,
2215 RegionReadState &readState,
2216 bool &isIsolatedFromAbove) {
2217 // Parse the name of the operation.
2218 std::optional<bool> wasRegistered;
2219 FailureOr<OperationName> opName = parseOpName(reader, wasRegistered);
2220 if (failed(Result: opName))
2221 return failure();
2222
2223 // Parse the operation mask, which indicates which components of the operation
2224 // are present.
2225 uint8_t opMask;
2226 if (failed(Result: reader.parseByte(value&: opMask)))
2227 return failure();
2228
2229 /// Parse the location.
2230 LocationAttr opLoc;
2231 if (failed(Result: parseAttribute(reader, result&: opLoc)))
2232 return failure();
2233
2234 // With the location and name resolved, we can start building the operation
2235 // state.
2236 OperationState opState(opLoc, *opName);
2237
2238 // Parse the attributes of the operation.
2239 if (opMask & bytecode::OpEncodingMask::kHasAttrs) {
2240 DictionaryAttr dictAttr;
2241 if (failed(parseAttribute(reader, dictAttr)))
2242 return failure();
2243 opState.attributes = dictAttr;
2244 }
2245
2246 if (opMask & bytecode::OpEncodingMask::kHasProperties) {
2247 // kHasProperties wasn't emitted in older bytecode, we should never get
2248 // there without also having the `wasRegistered` flag available.
2249 if (!wasRegistered)
2250 return emitError(loc: fileLoc,
2251 message: "Unexpected missing `wasRegistered` opname flag at "
2252 "bytecode version ")
2253 << version << " with properties.";
2254 // When an operation is emitted without being registered, the properties are
2255 // stored as an attribute. Otherwise the op must implement the bytecode
2256 // interface and control the serialization.
2257 if (wasRegistered) {
2258 DialectReader dialectReader(attrTypeReader, stringReader, resourceReader,
2259 dialectsMap, reader, version);
2260 if (failed(
2261 Result: propertiesReader.read(fileLoc, dialectReader, opName: &*opName, opState)))
2262 return failure();
2263 } else {
2264 // If the operation wasn't registered when it was emitted, the properties
2265 // was serialized as an attribute.
2266 if (failed(Result: parseAttribute(reader, result&: opState.propertiesAttr)))
2267 return failure();
2268 }
2269 }
2270
2271 /// Parse the results of the operation.
2272 if (opMask & bytecode::OpEncodingMask::kHasResults) {
2273 uint64_t numResults;
2274 if (failed(Result: reader.parseVarInt(result&: numResults)))
2275 return failure();
2276 opState.types.resize(N: numResults);
2277 for (int i = 0, e = numResults; i < e; ++i)
2278 if (failed(Result: parseType(reader, result&: opState.types[i])))
2279 return failure();
2280 }
2281
2282 /// Parse the operands of the operation.
2283 if (opMask & bytecode::OpEncodingMask::kHasOperands) {
2284 uint64_t numOperands;
2285 if (failed(Result: reader.parseVarInt(result&: numOperands)))
2286 return failure();
2287 opState.operands.resize(N: numOperands);
2288 for (int i = 0, e = numOperands; i < e; ++i)
2289 if (!(opState.operands[i] = parseOperand(reader)))
2290 return failure();
2291 }
2292
2293 /// Parse the successors of the operation.
2294 if (opMask & bytecode::OpEncodingMask::kHasSuccessors) {
2295 uint64_t numSuccs;
2296 if (failed(Result: reader.parseVarInt(result&: numSuccs)))
2297 return failure();
2298 opState.successors.resize(N: numSuccs);
2299 for (int i = 0, e = numSuccs; i < e; ++i) {
2300 if (failed(Result: parseEntry(reader, entries&: readState.curBlocks, entry&: opState.successors[i],
2301 entryStr: "successor")))
2302 return failure();
2303 }
2304 }
2305
2306 /// Parse the use-list orders for the results of the operation. Use-list
2307 /// orders are available since version 3 of the bytecode.
2308 std::optional<UseListMapT> resultIdxToUseListMap = std::nullopt;
2309 if (version >= bytecode::kUseListOrdering &&
2310 (opMask & bytecode::OpEncodingMask::kHasUseListOrders)) {
2311 size_t numResults = opState.types.size();
2312 auto parseResult = parseUseListOrderForRange(reader, numResults);
2313 if (failed(Result: parseResult))
2314 return failure();
2315 resultIdxToUseListMap = std::move(*parseResult);
2316 }
2317
2318 /// Parse the regions of the operation.
2319 if (opMask & bytecode::OpEncodingMask::kHasInlineRegions) {
2320 uint64_t numRegions;
2321 if (failed(Result: reader.parseVarIntWithFlag(result&: numRegions, flag&: isIsolatedFromAbove)))
2322 return failure();
2323
2324 opState.regions.reserve(N: numRegions);
2325 for (int i = 0, e = numRegions; i < e; ++i)
2326 opState.regions.push_back(Elt: std::make_unique<Region>());
2327 }
2328
2329 // Create the operation at the back of the current block.
2330 Operation *op = Operation::create(state: opState);
2331 readState.curBlock->push_back(op);
2332
2333 // If the operation had results, update the value references. We don't need to
2334 // do this if the current value scope is empty. That is, the op was not
2335 // encoded within a parent region.
2336 if (readState.numValues && op->getNumResults() &&
2337 failed(Result: defineValues(reader, values: op->getResults())))
2338 return failure();
2339
2340 /// Store a map for every value that received a custom use-list order from the
2341 /// bytecode file.
2342 if (resultIdxToUseListMap.has_value()) {
2343 for (size_t idx = 0; idx < op->getNumResults(); idx++) {
2344 if (resultIdxToUseListMap->contains(Val: idx)) {
2345 valueToUseListMap.try_emplace(Key: op->getResult(idx).getAsOpaquePointer(),
2346 Args: resultIdxToUseListMap->at(Val: idx));
2347 }
2348 }
2349 }
2350 return op;
2351}
2352
2353LogicalResult BytecodeReader::Impl::parseRegion(RegionReadState &readState) {
2354 EncodingReader &reader = *readState.reader;
2355
2356 // Parse the number of blocks in the region.
2357 uint64_t numBlocks;
2358 if (failed(Result: reader.parseVarInt(result&: numBlocks)))
2359 return failure();
2360
2361 // If the region is empty, there is nothing else to do.
2362 if (numBlocks == 0)
2363 return success();
2364
2365 // Parse the number of values defined in this region.
2366 uint64_t numValues;
2367 if (failed(Result: reader.parseVarInt(result&: numValues)))
2368 return failure();
2369 readState.numValues = numValues;
2370
2371 // Create the blocks within this region. We do this before processing so that
2372 // we can rely on the blocks existing when creating operations.
2373 readState.curBlocks.clear();
2374 readState.curBlocks.reserve(N: numBlocks);
2375 for (uint64_t i = 0; i < numBlocks; ++i) {
2376 readState.curBlocks.push_back(Elt: new Block());
2377 readState.curRegion->push_back(block: readState.curBlocks.back());
2378 }
2379
2380 // Prepare the current value scope for this region.
2381 valueScopes.back().push(readState);
2382
2383 // Parse the entry block of the region.
2384 readState.curBlock = readState.curRegion->begin();
2385 return parseBlockHeader(reader, readState);
2386}
2387
2388LogicalResult
2389BytecodeReader::Impl::parseBlockHeader(EncodingReader &reader,
2390 RegionReadState &readState) {
2391 bool hasArgs;
2392 if (failed(Result: reader.parseVarIntWithFlag(result&: readState.numOpsRemaining, flag&: hasArgs)))
2393 return failure();
2394
2395 // Parse the arguments of the block.
2396 if (hasArgs && failed(Result: parseBlockArguments(reader, block: &*readState.curBlock)))
2397 return failure();
2398
2399 // Uselist orders are available since version 3 of the bytecode.
2400 if (version < bytecode::kUseListOrdering)
2401 return success();
2402
2403 uint8_t hasUseListOrders = 0;
2404 if (hasArgs && failed(Result: reader.parseByte(value&: hasUseListOrders)))
2405 return failure();
2406
2407 if (!hasUseListOrders)
2408 return success();
2409
2410 Block &blk = *readState.curBlock;
2411 auto argIdxToUseListMap =
2412 parseUseListOrderForRange(reader, numResults: blk.getNumArguments());
2413 if (failed(Result: argIdxToUseListMap) || argIdxToUseListMap->empty())
2414 return failure();
2415
2416 for (size_t idx = 0; idx < blk.getNumArguments(); idx++)
2417 if (argIdxToUseListMap->contains(Val: idx))
2418 valueToUseListMap.try_emplace(Key: blk.getArgument(i: idx).getAsOpaquePointer(),
2419 Args: argIdxToUseListMap->at(Val: idx));
2420
2421 // We don't parse the operations of the block here, that's done elsewhere.
2422 return success();
2423}
2424
2425LogicalResult BytecodeReader::Impl::parseBlockArguments(EncodingReader &reader,
2426 Block *block) {
2427 // Parse the value ID for the first argument, and the number of arguments.
2428 uint64_t numArgs;
2429 if (failed(Result: reader.parseVarInt(result&: numArgs)))
2430 return failure();
2431
2432 SmallVector<Type> argTypes;
2433 SmallVector<Location> argLocs;
2434 argTypes.reserve(N: numArgs);
2435 argLocs.reserve(N: numArgs);
2436
2437 Location unknownLoc = UnknownLoc::get(config.getContext());
2438 while (numArgs--) {
2439 Type argType;
2440 LocationAttr argLoc = unknownLoc;
2441 if (version >= bytecode::kElideUnknownBlockArgLocation) {
2442 // Parse the type with hasLoc flag to determine if it has type.
2443 uint64_t typeIdx;
2444 bool hasLoc;
2445 if (failed(Result: reader.parseVarIntWithFlag(result&: typeIdx, flag&: hasLoc)) ||
2446 !(argType = attrTypeReader.resolveType(index: typeIdx)))
2447 return failure();
2448 if (hasLoc && failed(Result: parseAttribute(reader, result&: argLoc)))
2449 return failure();
2450 } else {
2451 // All args has type and location.
2452 if (failed(Result: parseType(reader, result&: argType)) ||
2453 failed(Result: parseAttribute(reader, result&: argLoc)))
2454 return failure();
2455 }
2456 argTypes.push_back(Elt: argType);
2457 argLocs.push_back(Elt: argLoc);
2458 }
2459 block->addArguments(types: argTypes, locs: argLocs);
2460 return defineValues(reader, values: block->getArguments());
2461}
2462
2463//===----------------------------------------------------------------------===//
2464// Value Processing
2465//===----------------------------------------------------------------------===//
2466
2467Value BytecodeReader::Impl::parseOperand(EncodingReader &reader) {
2468 std::vector<Value> &values = valueScopes.back().values;
2469 Value *value = nullptr;
2470 if (failed(Result: parseEntry(reader, entries&: values, entry&: value, entryStr: "value")))
2471 return Value();
2472
2473 // Create a new forward reference if necessary.
2474 if (!*value)
2475 *value = createForwardRef();
2476 return *value;
2477}
2478
2479LogicalResult BytecodeReader::Impl::defineValues(EncodingReader &reader,
2480 ValueRange newValues) {
2481 ValueScope &valueScope = valueScopes.back();
2482 std::vector<Value> &values = valueScope.values;
2483
2484 unsigned &valueID = valueScope.nextValueIDs.back();
2485 unsigned valueIDEnd = valueID + newValues.size();
2486 if (valueIDEnd > values.size()) {
2487 return reader.emitError(
2488 args: "value index range was outside of the expected range for "
2489 "the parent region, got [",
2490 args&: valueID, args: ", ", args&: valueIDEnd, args: "), but the maximum index was ",
2491 args: values.size() - 1);
2492 }
2493
2494 // Assign the values and update any forward references.
2495 for (unsigned i = 0, e = newValues.size(); i != e; ++i, ++valueID) {
2496 Value newValue = newValues[i];
2497
2498 // Check to see if a definition for this value already exists.
2499 if (Value oldValue = std::exchange(obj&: values[valueID], new_val&: newValue)) {
2500 Operation *forwardRefOp = oldValue.getDefiningOp();
2501
2502 // Assert that this is a forward reference operation. Given how we compute
2503 // definition ids (incrementally as we parse), it shouldn't be possible
2504 // for the value to be defined any other way.
2505 assert(forwardRefOp && forwardRefOp->getBlock() == &forwardRefOps &&
2506 "value index was already defined?");
2507
2508 oldValue.replaceAllUsesWith(newValue);
2509 forwardRefOp->moveBefore(block: &openForwardRefOps, iterator: openForwardRefOps.end());
2510 }
2511 }
2512 return success();
2513}
2514
2515Value BytecodeReader::Impl::createForwardRef() {
2516 // Check for an available existing operation to use. Otherwise, create a new
2517 // fake operation to use for the reference.
2518 if (!openForwardRefOps.empty()) {
2519 Operation *op = &openForwardRefOps.back();
2520 op->moveBefore(block: &forwardRefOps, iterator: forwardRefOps.end());
2521 } else {
2522 forwardRefOps.push_back(op: Operation::create(state: forwardRefOpState));
2523 }
2524 return forwardRefOps.back().getResult(idx: 0);
2525}
2526
2527//===----------------------------------------------------------------------===//
2528// Entry Points
2529//===----------------------------------------------------------------------===//
2530
2531BytecodeReader::~BytecodeReader() { assert(getNumOpsToMaterialize() == 0); }
2532
2533BytecodeReader::BytecodeReader(
2534 llvm::MemoryBufferRef buffer, const ParserConfig &config, bool lazyLoading,
2535 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
2536 Location sourceFileLoc =
2537 FileLineColLoc::get(context: config.getContext(), fileName: buffer.getBufferIdentifier(),
2538 /*line=*/0, /*column=*/0);
2539 impl = std::make_unique<Impl>(args&: sourceFileLoc, args: config, args&: lazyLoading, args&: buffer,
2540 args: bufferOwnerRef);
2541}
2542
2543LogicalResult BytecodeReader::readTopLevel(
2544 Block *block, llvm::function_ref<bool(Operation *)> lazyOpsCallback) {
2545 return impl->read(block, lazyOpsCallback);
2546}
2547
2548int64_t BytecodeReader::getNumOpsToMaterialize() const {
2549 return impl->getNumOpsToMaterialize();
2550}
2551
2552bool BytecodeReader::isMaterializable(Operation *op) {
2553 return impl->isMaterializable(op);
2554}
2555
2556LogicalResult BytecodeReader::materialize(
2557 Operation *op, llvm::function_ref<bool(Operation *)> lazyOpsCallback) {
2558 return impl->materialize(op, lazyOpsCallback);
2559}
2560
2561LogicalResult
2562BytecodeReader::finalize(function_ref<bool(Operation *)> shouldMaterialize) {
2563 return impl->finalize(shouldMaterialize);
2564}
2565
2566bool mlir::isBytecode(llvm::MemoryBufferRef buffer) {
2567 return buffer.getBuffer().starts_with(Prefix: "ML\xefR");
2568}
2569
2570/// Read the bytecode from the provided memory buffer reference.
2571/// `bufferOwnerRef` if provided is the owning source manager for the buffer,
2572/// and may be used to extend the lifetime of the buffer.
2573static LogicalResult
2574readBytecodeFileImpl(llvm::MemoryBufferRef buffer, Block *block,
2575 const ParserConfig &config,
2576 const std::shared_ptr<llvm::SourceMgr> &bufferOwnerRef) {
2577 Location sourceFileLoc =
2578 FileLineColLoc::get(context: config.getContext(), fileName: buffer.getBufferIdentifier(),
2579 /*line=*/0, /*column=*/0);
2580 if (!isBytecode(buffer)) {
2581 return emitError(loc: sourceFileLoc,
2582 message: "input buffer is not an MLIR bytecode file");
2583 }
2584
2585 BytecodeReader::Impl reader(sourceFileLoc, config, /*lazyLoading=*/false,
2586 buffer, bufferOwnerRef);
2587 return reader.read(block, /*lazyOpsCallback=*/nullptr);
2588}
2589
2590LogicalResult mlir::readBytecodeFile(llvm::MemoryBufferRef buffer, Block *block,
2591 const ParserConfig &config) {
2592 return readBytecodeFileImpl(buffer, block, config, /*bufferOwnerRef=*/{});
2593}
2594LogicalResult
2595mlir::readBytecodeFile(const std::shared_ptr<llvm::SourceMgr> &sourceMgr,
2596 Block *block, const ParserConfig &config) {
2597 return readBytecodeFileImpl(
2598 buffer: *sourceMgr->getMemoryBuffer(i: sourceMgr->getMainFileID()), block, config,
2599 bufferOwnerRef: sourceMgr);
2600}
2601

Provided by KDAB

Privacy Policy
Improve your Profiling and Debugging skills
Find out more

source code of mlir/lib/Bytecode/Reader/BytecodeReader.cpp