Add API and tool to process data files suitable for side storage of large attributes.

mlir/include/mlir/Support/DataFile.h

+//===- DataFile.h - I/O support for MLIR Data Files -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// A number of built-in MLIR attributes can represent arbitrarily large data
+// which is not suitable for direct inclusion in a module. There are various
+// mechanisms for externalizing this data, and hooks exist for projects to
+// bring their own serialization formats. This file, and the corresponding
+// mlir-data-file tool provide a basic implementation of a binary file format
+// that should be appropriate for many such use cases.
+//
+// Logically, an MLIR Data File presents in either read or write mode. When
+// reading, it maps an arbitrary symbol name to some binary contents. When
+// writing, the Data File can have duplicate symbols appended, with the last
+// one winning.
+//
+// Physically, an MLIR Data File consists of:
+//
+// ```
+// (File_Header_Record)
+// (
+//   (Content_Record)*
+//   (TOC_Record)*
+//   (Section_Trailer_Record)
+// )*
+// ```
+//
+// All records are minimally aligned to 16 byte offsets within the physical
+// file, and Content_Records can be individually aligned to larger values.
+//
+// Where the file format uses integers, they are encoded as little endian.
+//
+// Since Content_Records are opaque and referred to by physical offset in the
+// TOC, the Data File can be used in contexts disconnected from MLIR (i.e.
+// as an auxillary file to some other assembly that just references contents
+// by offset), and such usage can be optimized by first compacting and stripping
+// TOC records.
+//
+// Data files contain a file format version which must be supported by a
+// corresponding API that wishes to manipulate file metadata. This API only
+// supports writing to a data file at the same version as itself. Reading
+// historic versions should be supported.
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_DATA_FILE_H_
+#define MLIR_SUPPORT_DATA_FILE_H_
+
+#include "llvm/Support/Error.h"
+
+namespace mlir {
+
+class DataFileBuilder {
+public:
+};
+
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_DATA_FILE_H_

mlir/lib/Support/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
+  DataFile.cpp
   DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   InterfaceSupport.cpp
   StorageUniquer.cpp
   Timing.cpp
   ToolUtilities.cpp
   TypeID.cpp
 )
 
 add_mlir_library(MLIRSupport
+  DataFile.cpp
   DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   InterfaceSupport.cpp
   StorageUniquer.cpp
   Timing.cpp
   ToolUtilities.cpp
   TypeID.cpp

mlir/lib/Support/DataFile.cpp

+//===- DataFile.cpp - MLIR Data File Library ------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/DataFile.h"
+
+using namespace mlir;

mlir/tools/CMakeLists.txt

+add_subdirectory(mlir-data-file)
 add_subdirectory(mlir-lsp-server)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-parser-fuzzer)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-shlib)
 add_subdirectory(mlir-spirv-cpu-runner)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
 
 # mlir-cpu-runner requires ExecutionEngine which is only built
 # when the native target is configured in.
 if(TARGET ${LLVM_NATIVE_ARCH})
   add_subdirectory(mlir-cpu-runner)
 endif()

mlir/tools/mlir-data-file/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  Support
+  )
+
+add_llvm_tool(mlir-data-file mlir-data-file.cpp)
+llvm_update_compile_flags(mlir-data-file)
+target_link_libraries(mlir-data-file PRIVATE
+  MLIRSupport
+  MLIRParser
+)

mlir/tools/mlir-data-file/mlir-data-file.cpp

+//===- mlir-linalg-ods-yaml-gen.cpp - Linalg ODS generation from yaml  ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Operates on MLIR DataFiles. This tool doubles as both a testing tool of
+// the format and utilities and as a maintenance tool (i.e. compacting,
+// stripping, etc).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/DataFile.h"
+#include "llvm/Support/BinaryByteStream.h"
+#include "llvm/Support/BinaryStreamWriter.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileOutputBuffer.h"
+#include "llvm/Support/FileSystem.h"
+
+using namespace llvm;
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// DataFileBinaryStream
+// The stock implementation of DataFileBinaryStream assumes a fixed size,
+// write once modality. When building a Data File, we are operating on
+// a backing file of unknown length and we want to operate on the file
+// directly.
+// TODO: Split this into a read-only base class and a writable subclass.
+//===----------------------------------------------------------------------===//
+
+class DataFileBinaryStream : public WritableBinaryStream {
+public:
+  ~DataFileBinaryStream() override;
+
+  static Expected<std::unique_ptr<DataFileBinaryStream>>
+  openForReadWrite(StringRef filePath);
+
+  llvm::support::endianness getEndian() const override {
+    return llvm::support::endianness::little;
+  }
+
+  Error readBytes(uint64_t Offset, uint64_t Size,
+                  ArrayRef<uint8_t> &Buffer) override;
+
+  Error readLongestContiguousChunk(uint64_t Offset,
+                                   ArrayRef<uint8_t> &Buffer) override;
+
+  uint64_t getLength() override;
+
+  BinaryStreamFlags getFlags() const override {
+    if (writable) {
+      return BSF_Write | BSF_Append;
+    } else {
+      return BSF_None;
+    }
+  }
+
+  Error writeBytes(uint64_t Offset, ArrayRef<uint8_t> Data) override;
+
+  Error commit() override;
+
+private:
+  DataFileBinaryStream(llvm::sys::fs::file_t file, uint64_t fileSize,
+                       bool writable)
+      : file(file), os(file, /*shouldClose=*/false, /*unbuffered=*/true),
+        fileSize(fileSize), writable(writable) {
+    // Since we don't use the read buffer for bulk data-access, we can be
+    // fairly small here.
+    readBuffer.resize(4096);
+  }
+  llvm::sys::fs::file_t file;
+  raw_fd_ostream os;
+  // We fulfill read requests from this buffer.
+  std::vector<uint8_t> readBuffer;
+  uint64_t fileSize;
+  bool writable;
+};
+
+Expected<std::unique_ptr<DataFileBinaryStream>>
+DataFileBinaryStream::openForReadWrite(StringRef filePath) {
+  // Open file.
+  auto file = llvm::sys::fs::openNativeFileForReadWrite(
+      filePath, llvm::sys::fs::CD_OpenAlways, llvm::sys::fs::OF_None);
+  if (auto E = file.takeError()) {
+    return std::move(E);
+  }
+
+  // Get its current size.
+  llvm::sys::fs::file_status fileStatus;
+  if (auto ec = llvm::sys::fs::status(*file, fileStatus)) {
+    llvm::sys::fs::closeFile(*file);
+    return errorCodeToError(ec);
+  }
+
+  return std::unique_ptr<DataFileBinaryStream>(
+      new DataFileBinaryStream(std::move(*file), fileStatus.getSize(), true));
+}
+
+DataFileBinaryStream::~DataFileBinaryStream() {
+  llvm::sys::fs::closeFile(file);
+}
+
+Error DataFileBinaryStream::commit() { return Error::success(); }
+
+uint64_t DataFileBinaryStream::getLength() { return fileSize; }
+
+Error DataFileBinaryStream::readBytes(uint64_t Offset, uint64_t Size,
+                                      ArrayRef<uint8_t> &Buffer) {
+  auto actualSize = llvm::sys::fs::readNativeFileSlice(
+      file,
+      MutableArrayRef<char>(reinterpret_cast<char *>(readBuffer.data()),
+                            readBuffer.size()),
+      Offset);
+  if (auto E = actualSize.takeError()) {
+    return E;
+  }
+  Buffer = ArrayRef<uint8_t>(readBuffer.data(), *actualSize);
+  return Error::success();
+}
+
+Error DataFileBinaryStream::readLongestContiguousChunk(
+    uint64_t Offset, ArrayRef<uint8_t> &Buffer) {
+  return readBytes(Offset, readBuffer.size(), Buffer);
+}
+
+Error DataFileBinaryStream::writeBytes(uint64_t Offset,
+                                       ArrayRef<uint8_t> Data) {
+  os.pwrite(reinterpret_cast<const char *>(Data.data()), Data.size(), Offset);
+  auto ec = os.error();
+  if (ec) {
+    os.clear_error();
+    return errorCodeToError(ec);
+  }
+
+  fileSize = std::max(fileSize, Offset + Data.size());
+  return Error::success();
+}
+
+//===----------------------------------------------------------------------===//
+// Create SubCommand
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct CreateSubCommand {
+  CreateSubCommand()
+      : subCommand("create", "Creates an MLIR Data File from the command line"),
+        outputFile(cl::Positional, cl::sub(subCommand),
+                   cl::desc("<file to create>"), cl::Required),
+        specs(cl::Positional, cl::sub(subCommand),
+              cl::desc("(<symbol> <value>)*"), cl::ZeroOrMore) {}
+  cl::SubCommand subCommand;
+  cl::opt<std::string> outputFile;
+  cl::list<std::string> specs;
+
+  int run();
+};
+} // namespace
+
+int CreateSubCommand::run() {
+  // Open file.
+  auto expectedStream = DataFileBinaryStream::openForReadWrite(outputFile);
+  if (auto E = expectedStream.takeError()) {
+    errs() << "Could not create output file (" << outputFile
+           << "): " << toString(std::move(E)) << "\n";
+    return 1;
+  }
+
+  // Access stream.
+  std::unique_ptr<DataFileBinaryStream> stream = std::move(*expectedStream);
+  BinaryStreamWriter writer(*stream);
+  if (auto E = writer.writeCString("Foobar!\n")) {
+    errs() << "Error writing: " << toString(std::move(E)) << "\n";
+    return 1;
+  }
+
+  if (auto E = writer.padToAlignment(64)) {
+    errs() << "Error padding: " << toString(std::move(E)) << "\n";
+    return 1;
+  }
+
+  if (auto E = writer.writeCString("Still here!\n")) {
+    errs() << "Error writing: " << toString(std::move(E)) << "\n";
+    return 1;
+  }
+
+  // Commit.
+  if (auto E = stream->commit()) {
+    errs() << "Error committing output file: " << toString(std::move(E))
+           << "\n";
+    return 1;
+  }
+
+  return 0;
+}
+
+int main(int argc, char **argv) {
+  static CreateSubCommand createSc;
+
+  cl::ParseCommandLineOptions(argc, argv, "MLIR Data File Tool");
+
+  if (createSc.subCommand) {
+    return createSc.run();
+  } else {
+    llvm::errs() << "Unknown subcommand\n";
+    return 1;
+  }
+
+  return 0;
+}