Method getBytecodeSize

lib/Bytecode/Reader/BytecodeReader.cpp:2416–2459

//===----------------------------------------------------------------------===//

std::optional<size_t> cuda_tile::getBytecodeSize(const char *bytecodeBuffer) {
  if (!isTileIRBytecode(bytecodeBuffer))
    return std::nullopt;

  auto charBuffer = reinterpret_cast<const unsigned char*>(bytecodeBuffer);
  if (charBuffer[sizeof(kTileIRBytecodeMagic)] == 0)
    return std::nullopt;

  // Build a buffer assuming we have the maximum size of the bytecode, we'll
  // infer the actual size as we parse the bytecode.
  ArrayRef<uint8_t> bytecodeData(
      reinterpret_cast<const uint8_t *>(bytecodeBuffer),
      SIZE_MAX - reinterpret_cast<uintptr_t>(bytecodeBuffer));

  // Set up the reader and context.
  MLIRContext context(MLIRContext::Threading::DISABLED);
  EncodingReader reader(bytecodeData, context);
  ScopedDiagnosticHandler handler(&context, [](Diagnostic &diag) {
    // Ignore all errors.
  });

  // Parse the header of the bytecode.
  BytecodeVersion version;
  if (failed(parseHeader(reader, context, version)))
    return std::nullopt;

  // Parse the sections until we reach the end of the bytecode. We don't
  // actually try to reason about the section data, we just want to know the
  // sizes.
  std::array<bool, Section::NumSections> seenSections;
  seenSections.fill(false);
  while (true) {
    // Parse the next section.
    SectionHeader header;
    if (failed(parseSectionHeader(reader, header, context)) ||
        failed(reader.skip(header.length)) ||
        std::exchange(seenSections[header.sectionID], true))
      return std::nullopt;

    // Check for the end of the bytecode stream.
    if (header.sectionID == Section::EndOfBytecode)
      return reader.currentOffset();
  }
}

OwningOpRef<cuda_tile::ModuleOp>
cuda_tile::readBytecode(llvm::MemoryBufferRef bytecodeBuffer,