Adding IREE parameter archive format and tooling support. (#15670)
The new format allows us to store parameters for both inference and
training observing the requirements for both efficient CPU and
GPU/accelerator execution. We can also support additional storage types
such as splats allowing for stripped parameter files that work with
programs compiled assuming real parameters. The format has a provision
for referencing external file ranges but support for reading such files
is TBD.
The format can be used like gguf/safetensors and is supported by the
tooling in the same way. Additionally a new `iree-convert-parameters`
tool is added to convert any format supported for loading (gguf,
safetensors, and irpa itself) into irpa files with some control over
which parameters are included, renaming of parameters, and stripping
parameters and replacing them with splat values. This should make it
easy to take any gguf/safetensors file and quickly create stripped
variants for easy reproducers/CI benchmarking without needing to ship
around the original files. The `iree-create-parameters` tool can be used
to create empty archives that are ready for initialization from a
program that mutates them or to create parameter archives with named
parameters when there is no source gguf/safetensors file.
All of this is still using memory-mapped files; this limits our
parameter file sizes on 32-bit systems but I suspect no one is going to
run this tool for large models on 32-bit systems. In the future we can
make the conversion tool use the HAL and schedule out optimized file
I/O. For now we just copy parameters via a normal read/write loop and
it's fastish-enough (pretty much I/O bound, with less optimal reads
because of memory mapping). For me with a cold cache it takes ~1min to
rewrite a 25GB file and 25sec with a hot cache.
This initial commit has the IRPA builder using the new iree_io_stream_t
but switching all format parsers to use it is deferred to future
changes.
Progress on #15521.
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt
index e8a7be4..e771549 100644
--- a/runtime/bindings/python/CMakeLists.txt
+++ b/runtime/bindings/python/CMakeLists.txt
@@ -89,6 +89,7 @@
iree::base::internal::path
iree::io::file_handle
iree::io::formats::gguf
+ iree::io::formats::irpa
iree::io::formats::safetensors
iree::io::parameter_index
iree::io::parameter_index_provider
@@ -140,6 +141,8 @@
"iree/_runtime/scripts/iree_benchmark_trace/__main__.py"
"iree/_runtime/scripts/iree_benchmark_module/__main__.py"
"iree/_runtime/scripts/iree_cpuinfo/__main__.py"
+ "iree/_runtime/scripts/iree_convert_parameters/__main__.py"
+ "iree/_runtime/scripts/iree_create_parameters/__main__.py"
"iree/_runtime/scripts/iree_dump_module/__main__.py"
"iree/_runtime/scripts/iree_dump_parameters/__main__.py"
"iree/_runtime/scripts/iree_run_trace/__main__.py"
@@ -163,12 +166,24 @@
iree_symlink_tool(
TARGET runtime
+ FROM_TOOL_TARGET iree-convert-parameters
+ TO_EXE_NAME iree/_runtime_libs/iree-convert-parameters
+)
+
+iree_symlink_tool(
+ TARGET runtime
FROM_TOOL_TARGET iree-cpuinfo
TO_EXE_NAME iree/_runtime_libs/iree-cpuinfo
)
iree_symlink_tool(
TARGET runtime
+ FROM_TOOL_TARGET iree-create-parameters
+ TO_EXE_NAME iree/_runtime_libs/iree-create-parameters
+)
+
+iree_symlink_tool(
+ TARGET runtime
FROM_TOOL_TARGET iree-dump-module
TO_EXE_NAME iree/_runtime_libs/iree-dump-module
)
@@ -315,6 +330,8 @@
iree-cpuinfo
iree-benchmark-module
iree-benchmark-trace
+ iree-convert-parameters
+ iree-create-parameters
iree-dump-module
iree-dump-parameters
iree-run-module
diff --git a/runtime/bindings/python/io.cc b/runtime/bindings/python/io.cc
index 0d8f8d1..c2a30a4 100644
--- a/runtime/bindings/python/io.cc
+++ b/runtime/bindings/python/io.cc
@@ -14,8 +14,9 @@
#include "./vm.h"
#include "iree/base/internal/file_io.h"
#include "iree/base/internal/path.h"
-#include "iree/io/formats/gguf/gguf_format.h"
-#include "iree/io/formats/safetensors/safetensors_format.h"
+#include "iree/io/formats/gguf/gguf_parser.h"
+#include "iree/io/formats/irpa/irpa_parser.h"
+#include "iree/io/formats/safetensors/safetensors_parser.h"
#include "iree/io/parameter_index_provider.h"
#include "iree/modules/io/parameters/module.h"
@@ -101,13 +102,18 @@
CheckApiStatus(
iree_io_parse_gguf_index(file_handle.raw_ptr(), self.raw_ptr()),
"Could not parse gguf file into index");
+ } else if (format == "irpa") {
+ CheckApiStatus(
+ iree_io_parse_irpa_index(file_handle.raw_ptr(), self.raw_ptr()),
+ "Could not parse IREE parameter archive file into index");
} else if (format == "safetensors") {
CheckApiStatus(
iree_io_parse_safetensors_index(file_handle.raw_ptr(), self.raw_ptr()),
"Could not parse safetensors file into index");
} else {
throw std::invalid_argument(
- "Unrecognized file format. Expected one of: 'gguf', 'safetensors'");
+ "Unrecognized file format. Expected one of: 'gguf', 'irpa', "
+ "'safetensors'");
}
}
diff --git a/runtime/bindings/python/iree/_runtime/scripts/iree_convert_parameters/__main__.py b/runtime/bindings/python/iree/_runtime/scripts/iree_convert_parameters/__main__.py
new file mode 100644
index 0000000..9e71002
--- /dev/null
+++ b/runtime/bindings/python/iree/_runtime/scripts/iree_convert_parameters/__main__.py
@@ -0,0 +1,23 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+import os
+import subprocess
+import sys
+
+# Note that iree-convert-parameters is only in the default libs.
+from iree import _runtime_libs
+
+
+def main(args=None):
+ if args is None:
+ args = sys.argv[1:]
+ exe = os.path.join(_runtime_libs.__path__[0], "iree-convert-parameters")
+ return subprocess.call(args=[exe] + args)
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/runtime/bindings/python/iree/_runtime/scripts/iree_create_parameters/__main__.py b/runtime/bindings/python/iree/_runtime/scripts/iree_create_parameters/__main__.py
new file mode 100644
index 0000000..09006cc
--- /dev/null
+++ b/runtime/bindings/python/iree/_runtime/scripts/iree_create_parameters/__main__.py
@@ -0,0 +1,23 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+import os
+import subprocess
+import sys
+
+# Note that iree-create-parameters is only in the default libs.
+from iree import _runtime_libs
+
+
+def main(args=None):
+ if args is None:
+ args = sys.argv[1:]
+ exe = os.path.join(_runtime_libs.__path__[0], "iree-create-parameters")
+ return subprocess.call(args=[exe] + args)
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/runtime/bindings/python/iree/runtime/_package_test.py b/runtime/bindings/python/iree/runtime/_package_test.py
index 1b51819..0f7124c 100644
--- a/runtime/bindings/python/iree/runtime/_package_test.py
+++ b/runtime/bindings/python/iree/runtime/_package_test.py
@@ -38,6 +38,8 @@
check_tool("iree-benchmark-trace", ["--help"], "IREE: iree-benchmark-trace")
check_tool("iree-run-module", ["--help"], "IREE: iree-run-module")
check_tool("iree-run-trace", ["--help"], "IREE: iree-run-trace")
+check_tool("iree-create-parameters", ["--help"], "IREE: iree-create-parameters")
+check_tool("iree-convert-parameters", ["--help"], "IREE: iree-convert-parameters")
check_tool("iree-dump-module", ["--help"], "IREE: iree-dump-module")
check_tool("iree-dump-parameters", ["--help"], "IREE: iree-dump-parameters")
check_tool("iree-cpuinfo", [])
diff --git a/runtime/bindings/python/tests/testdata/generate_parameter_gguf.py b/runtime/bindings/python/tests/testdata/generate_parameter_gguf.py
index 607f57e..0c518b9 100755
--- a/runtime/bindings/python/tests/testdata/generate_parameter_gguf.py
+++ b/runtime/bindings/python/tests/testdata/generate_parameter_gguf.py
@@ -22,10 +22,6 @@
writer.add_architecture()
writer.add_custom_alignment(64)
- writer.add_uint32("metadata_uint32", 42)
- writer.add_string("metadata_str", "hello")
- writer.add_array("metadata_strs", ["a", "b", "c"])
-
for key, value in tensors.items():
writer.add_tensor(key, value)
diff --git a/runtime/setup.py b/runtime/setup.py
index b4c8950..e561c45 100644
--- a/runtime/setup.py
+++ b/runtime/setup.py
@@ -569,6 +569,8 @@
"iree-benchmark-module*",
"iree-benchmark-trace*",
# These utilities are invariant wrt tracing and are only built for the default runtime.
+ "iree-create-parameters*",
+ "iree-convert-parameters*",
"iree-dump-module*",
"iree-dump-parameters*",
"iree-cpuinfo*",
@@ -593,6 +595,8 @@
"iree-run-trace = iree._runtime.scripts.iree_run_trace.__main__:main",
"iree-benchmark-module = iree._runtime.scripts.iree_benchmark_module.__main__:main",
"iree-benchmark-trace = iree._runtime.scripts.iree_benchmark_trace.__main__:main",
+ "iree-create-parameters = iree._runtime.scripts.iree_create_parameters.__main__:main",
+ "iree-convert-parameters = iree._runtime.scripts.iree_convert_parameters.__main__:main",
"iree-dump-module = iree._runtime.scripts.iree_dump_module.__main__:main",
"iree-dump-parameters = iree._runtime.scripts.iree_dump_parameters.__main__:main",
"iree-cpuinfo = iree._runtime.scripts.iree_cpuinfo.__main__:main",
diff --git a/runtime/src/iree/base/internal/file_io.c b/runtime/src/iree/base/internal/file_io.c
index f0e3804..2f61945 100644
--- a/runtime/src/iree/base/internal/file_io.c
+++ b/runtime/src/iree/base/internal/file_io.c
@@ -275,6 +275,41 @@
return iree_ok_status();
}
+static iree_status_t iree_file_create_mapped_platform(
+ const char* path, uint64_t file_size, uint64_t offset,
+ iree_host_size_t length, iree_file_contents_t* contents) {
+ // Create file on disk.
+ FILE* file = fopen(path, "w+b");
+ if (file == NULL) {
+ return iree_make_status(IREE_STATUS_NOT_FOUND, "failed to open file '%s'",
+ path);
+ }
+ contents->mapping = file;
+
+ // Zero-extend the file ('truncate' can extend, because... unix).
+ if (ftruncate(fileno(file), (off_t)file_size) == -1) {
+ return iree_make_status(iree_status_code_from_errno(errno),
+ "failed to extend file '%s' to %" PRIu64
+ " bytes (out of disk space or permission denied)",
+ path, file_size);
+ }
+
+ // Map the memory.
+ void* ptr = mmap(NULL, (size_t)length, PROT_READ | PROT_WRITE, MAP_SHARED,
+ fileno(file), (off_t)offset);
+ if (ptr == MAP_FAILED) {
+ return iree_make_status(iree_status_code_from_errno(errno),
+ "failed to map '%s' range %" PRIu64 "-%" PRIu64
+ " (%" PRIhsz " bytes) from file of %" PRIu64
+ " total bytes",
+ path, offset, offset + length, length, file_size);
+ }
+
+ contents->const_buffer =
+ iree_make_const_byte_span(ptr, (iree_host_size_t)length);
+ return iree_ok_status();
+}
+
static void iree_file_contents_free_platform(iree_file_contents_t* contents) {
if (contents->mapping) {
munmap(contents->buffer.data, (size_t)contents->buffer.data_length);
@@ -334,6 +369,67 @@
return iree_ok_status();
}
+static iree_status_t iree_file_create_mapped_platform(
+ const char* path, uint64_t file_size, uint64_t offset,
+ iree_host_size_t length, iree_file_contents_t* contents) {
+ // TODO(benvanik): investigate FILE_FLAG_SEQUENTIAL_SCAN +
+ // FILE_FLAG_WRITE_THROUGH flags once we have some benchmarks.
+ HANDLE file = CreateFileA(path, GENERIC_READ | GENERIC_WRITE, 0, NULL,
+ CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
+ if (file == INVALID_HANDLE_VALUE) {
+ return iree_make_status(iree_status_code_from_win32_error(GetLastError()),
+ "failed to open file '%s'", path);
+ }
+
+ // Zero-extend the file up to the total file size specified by the caller.
+ // This may be larger than the virtual address space can handle but so long as
+ // the length requested for mapping is under the size_t limit this will
+ // succeed.
+ LARGE_INTEGER file_size_li;
+ file_size_li.QuadPart = file_size;
+ if (!SetFilePointerEx(file, file_size_li, NULL, FILE_BEGIN) ||
+ !SetEndOfFile(file)) {
+ CloseHandle(file);
+ return iree_make_status(iree_status_code_from_win32_error(GetLastError()),
+ "failed to extend file '%s' to %" PRIu64
+ " bytes (out of disk space or permission denied)",
+ path, file_size);
+ }
+
+ // Create a file mapping object which will retain the file handle for the
+ // lifetime of the mapping.
+ HANDLE mapping =
+ CreateFileMappingA(file, NULL, PAGE_READWRITE, /*dwMaximumSizeHigh=*/0,
+ /*dwMaximumSizeLow=*/0, /*lpName=*/NULL);
+ if (!mapping) {
+ CloseHandle(file);
+ return iree_make_status(iree_status_code_from_win32_error(GetLastError()),
+ "failed to create file mapping for '%s'", path);
+ }
+
+ // Retained by the mapping so safe to release now.
+ CloseHandle(file);
+
+ // Map the requested range into the virtual address space of the process.
+ LARGE_INTEGER offset_li;
+ offset_li.QuadPart = offset;
+ void* ptr = MapViewOfFileEx(mapping, FILE_MAP_ALL_ACCESS, offset_li.HighPart,
+ offset_li.LowPart, (SIZE_T)length,
+ /*lpBaseAddress=*/NULL);
+ if (!ptr) {
+ CloseHandle(mapping);
+ return iree_make_status(iree_status_code_from_win32_error(GetLastError()),
+ "failed to map '%s' range %" PRIu64 "-%" PRIu64
+ " (%" PRIhsz " bytes) from file of %" PRIu64
+ " total bytes",
+ path, offset, offset + length, length, file_size);
+ }
+
+ contents->mapping = mapping;
+ contents->buffer = iree_make_byte_span(ptr, length);
+ return iree_ok_status();
+}
+
static void iree_file_contents_free_platform(iree_file_contents_t* contents) {
if (contents->mapping) {
UnmapViewOfFile(contents->buffer.data);
@@ -349,10 +445,42 @@
"file mapping not supported on this platform");
}
+static iree_status_t iree_file_create_mapped_platform(
+ const char* path, uint64_t file_size, uint64_t offset,
+ iree_host_size_t length, iree_file_contents_t* contents) {
+ return iree_make_status(IREE_STATUS_UNAVAILABLE,
+ "file mapping not supported on this platform");
+}
+
static void iree_file_contents_free_platform(iree_file_contents_t* contents) {}
#endif // IREE_PLATFORM_*
+iree_status_t iree_file_create_mapped(const char* path, uint64_t file_size,
+ uint64_t offset, iree_host_size_t length,
+ iree_allocator_t allocator,
+ iree_file_contents_t** out_contents) {
+ IREE_ASSERT_ARGUMENT(path);
+ IREE_ASSERT_ARGUMENT(out_contents);
+ *out_contents = NULL;
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_file_contents_t* contents = NULL;
+ iree_allocator_malloc(allocator, sizeof(*contents), (void**)&contents);
+ contents->allocator = allocator;
+
+ iree_status_t status = iree_file_create_mapped_platform(
+ path, file_size, offset, length, contents);
+
+ if (iree_status_is_ok(status)) {
+ *out_contents = contents;
+ } else {
+ iree_file_contents_free(contents);
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
iree_status_t iree_file_write_contents(const char* path,
iree_const_byte_span_t content) {
IREE_TRACE_ZONE_BEGIN(z0);
diff --git a/runtime/src/iree/base/internal/file_io.h b/runtime/src/iree/base/internal/file_io.h
index 581c2b9..7f64b58 100644
--- a/runtime/src/iree/base/internal/file_io.h
+++ b/runtime/src/iree/base/internal/file_io.h
@@ -88,6 +88,25 @@
const char* path, iree_allocator_t allocator,
iree_file_contents_t** out_contents);
+// EXPERIMENTAL: will be moved to iree_io_file_handle_t (along with the rest of
+// this file) in future changes.
+//
+// Creates a new file on disk and maps its contents for read/write access.
+// The file will be zero-extended up to |file_size| bytes and the returned
+// contents will start at |offset| and run for |length| bytes.
+// As pages may be read or written on demand care should be used when
+// profiling/benchmarking as warm-up costs will be higher and variance during
+// execution will go up. It's possible for dozens of GiB of outstanding writes
+// to complete at memcpy speed only to flush after the process has exited.
+//
+// Returns the contents of the file in |out_contents|.
+// |allocator| is used to allocate the memory and the caller must use
+// iree_file_contents_free to release the memory.
+iree_status_t iree_file_create_mapped(const char* path, uint64_t file_size,
+ uint64_t offset, iree_host_size_t length,
+ iree_allocator_t allocator,
+ iree_file_contents_t** out_contents);
+
// Synchronously writes a byte buffer into a file.
// Existing contents are overwritten.
iree_status_t iree_file_write_contents(const char* path,
diff --git a/runtime/src/iree/io/BUILD.bazel b/runtime/src/iree/io/BUILD.bazel
index 7885b8d..7e81715 100644
--- a/runtime/src/iree/io/BUILD.bazel
+++ b/runtime/src/iree/io/BUILD.bazel
@@ -21,6 +21,7 @@
"file_handle.h",
],
deps = [
+ ":stream",
"//runtime/src/iree/base",
"//runtime/src/iree/base/internal",
],
diff --git a/runtime/src/iree/io/CMakeLists.txt b/runtime/src/iree/io/CMakeLists.txt
index 95b18c3..af41f50 100644
--- a/runtime/src/iree/io/CMakeLists.txt
+++ b/runtime/src/iree/io/CMakeLists.txt
@@ -18,6 +18,7 @@
SRCS
"file_handle.c"
DEPS
+ ::stream
iree::base
iree::base::internal
PUBLIC
diff --git a/runtime/src/iree/io/file_handle.c b/runtime/src/iree/io/file_handle.c
index 97d1225..dc06b78 100644
--- a/runtime/src/iree/io/file_handle.c
+++ b/runtime/src/iree/io/file_handle.c
@@ -99,3 +99,112 @@
IREE_ASSERT_ARGUMENT(handle);
return handle->primitive;
}
+
+IREE_API_EXPORT iree_status_t
+iree_io_file_handle_flush(iree_io_file_handle_t* handle) {
+ IREE_ASSERT_ARGUMENT(handle);
+ IREE_TRACE_ZONE_BEGIN(z0);
+ iree_status_t status = iree_ok_status();
+ switch (handle->primitive.type) {
+ case IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION: {
+ // No-op (though we could flush when known mapped).
+ break;
+ }
+ default: {
+ status = iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "flush not supported on handle type %d",
+ (int)handle->primitive.type);
+ break;
+ }
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+//===----------------------------------------------------------------------===//
+// iree_io_stream_t utilities
+//===----------------------------------------------------------------------===//
+
+static void iree_io_memory_stream_file_release(void* user_data,
+ iree_io_stream_t* stream) {
+ iree_io_file_handle_t* file_handle = (iree_io_file_handle_t*)user_data;
+ iree_io_file_handle_release(file_handle);
+}
+
+IREE_API_EXPORT iree_status_t iree_io_stream_open(
+ iree_io_stream_mode_t mode, iree_io_file_handle_t* file_handle,
+ uint64_t file_offset, iree_allocator_t host_allocator,
+ iree_io_stream_t** out_stream) {
+ IREE_ASSERT_ARGUMENT(file_handle);
+ IREE_ASSERT_ARGUMENT(out_stream);
+ *out_stream = NULL;
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_status_t status = iree_ok_status();
+ iree_io_stream_t* stream = NULL;
+ iree_io_file_handle_primitive_t file_primitive =
+ iree_io_file_handle_primitive(file_handle);
+ switch (file_primitive.type) {
+ case IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION: {
+ if (file_offset > file_primitive.value.host_allocation.data_length) {
+ status = iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "file offset %" PRIu64
+ " out of range of host allocation with %" PRIhsz " bytes available",
+ file_offset, file_primitive.value.host_allocation.data_length);
+ break;
+ }
+ iree_io_memory_stream_release_callback_t release_callback = {
+ .fn = iree_io_memory_stream_file_release,
+ .user_data = file_handle,
+ };
+ iree_io_file_handle_retain(file_handle);
+ status = iree_io_memory_stream_wrap(
+ mode,
+ iree_make_byte_span(
+ file_primitive.value.host_allocation.data + file_offset,
+ file_primitive.value.host_allocation.data_length - file_offset),
+ release_callback, host_allocator, &stream);
+ if (!iree_status_is_ok(status)) iree_io_file_handle_release(file_handle);
+ break;
+ }
+ default: {
+ status =
+ iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "stream open not yet supported on handle type %d",
+ (int)file_primitive.type);
+ break;
+ }
+ }
+
+ if (iree_status_is_ok(status)) {
+ *out_stream = stream;
+ } else {
+ iree_io_stream_release(stream);
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+IREE_API_EXPORT iree_status_t iree_io_stream_write_file(
+ iree_io_stream_t* stream, iree_io_file_handle_t* source_file_handle,
+ uint64_t source_file_offset, iree_io_stream_pos_t length,
+ iree_allocator_t host_allocator) {
+ IREE_ASSERT_ARGUMENT(stream);
+ IREE_ASSERT_ARGUMENT(source_file_handle);
+ IREE_TRACE_ZONE_BEGIN(z0);
+ IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)length);
+
+ iree_io_stream_t* source_stream = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_stream_open(IREE_IO_STREAM_MODE_READABLE, source_file_handle,
+ source_file_offset, host_allocator, &source_stream));
+
+ iree_status_t status = iree_io_stream_copy(source_stream, stream, length);
+
+ iree_io_stream_release(source_stream);
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
diff --git a/runtime/src/iree/io/file_handle.h b/runtime/src/iree/io/file_handle.h
index 5b01a50..6279b69 100644
--- a/runtime/src/iree/io/file_handle.h
+++ b/runtime/src/iree/io/file_handle.h
@@ -10,6 +10,7 @@
#include <stdint.h>
#include "iree/base/api.h"
+#include "iree/io/stream.h"
#ifdef __cplusplus
extern "C" {
@@ -146,6 +147,29 @@
return iree_io_file_handle_primitive(handle).value;
}
+// Flushes pending writes of |handle| to its backing storage.
+IREE_API_EXPORT iree_status_t
+iree_io_file_handle_flush(iree_io_file_handle_t* handle);
+
+//===----------------------------------------------------------------------===//
+// iree_io_stream_t utilities
+//===----------------------------------------------------------------------===//
+
+// Opens a stream from the given |file_handle| at the absolute |file_offset|.
+// The returned stream will retain the file until it is released.
+IREE_API_EXPORT iree_status_t iree_io_stream_open(
+ iree_io_stream_mode_t mode, iree_io_file_handle_t* file_handle,
+ uint64_t file_offset, iree_allocator_t host_allocator,
+ iree_io_stream_t** out_stream);
+
+// Writes up to |length| bytes of |source_file_handle| starting at offset
+// |source_file_offset| to the target |stream|. |host_allocator| may be used
+// for transient allocations required during file I/O.
+IREE_API_EXPORT iree_status_t iree_io_stream_write_file(
+ iree_io_stream_t* stream, iree_io_file_handle_t* source_file_handle,
+ uint64_t source_file_offset, iree_io_stream_pos_t length,
+ iree_allocator_t host_allocator);
+
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
diff --git a/runtime/src/iree/io/formats/gguf/BUILD.bazel b/runtime/src/iree/io/formats/gguf/BUILD.bazel
index 879a34c..e0bb5a1 100644
--- a/runtime/src/iree/io/formats/gguf/BUILD.bazel
+++ b/runtime/src/iree/io/formats/gguf/BUILD.bazel
@@ -15,10 +15,10 @@
iree_runtime_cc_library(
name = "gguf",
srcs = [
- "gguf_format.c",
+ "gguf_parser.c",
],
hdrs = [
- "gguf_format.h",
+ "gguf_parser.h",
],
deps = [
"//runtime/src/iree/base",
@@ -28,12 +28,10 @@
)
iree_runtime_cc_test(
- name = "gguf_format_test",
- srcs = ["gguf_format_test.cc"],
- tags = ["requires-filesystem"],
+ name = "gguf_parser_test",
+ srcs = ["gguf_parser_test.cc"],
deps = [
":gguf",
- "//runtime/src/iree/base/internal:file_io",
"//runtime/src/iree/io/formats/gguf/testdata:gguf_files",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
diff --git a/runtime/src/iree/io/formats/gguf/CMakeLists.txt b/runtime/src/iree/io/formats/gguf/CMakeLists.txt
index c727aac..1b2b1de 100644
--- a/runtime/src/iree/io/formats/gguf/CMakeLists.txt
+++ b/runtime/src/iree/io/formats/gguf/CMakeLists.txt
@@ -14,9 +14,9 @@
NAME
gguf
HDRS
- "gguf_format.h"
+ "gguf_parser.h"
SRCS
- "gguf_format.c"
+ "gguf_parser.c"
DEPS
iree::base
iree::io::file_handle
@@ -26,17 +26,14 @@
iree_cc_test(
NAME
- gguf_format_test
+ gguf_parser_test
SRCS
- "gguf_format_test.cc"
+ "gguf_parser_test.cc"
DEPS
::gguf
- iree::base::internal::file_io
iree::io::formats::gguf::testdata::gguf_files
iree::testing::gtest
iree::testing::gtest_main
- LABELS
- "requires-filesystem"
)
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/io/formats/gguf/gguf_format.c b/runtime/src/iree/io/formats/gguf/gguf_parser.c
similarity index 99%
rename from runtime/src/iree/io/formats/gguf/gguf_format.c
rename to runtime/src/iree/io/formats/gguf/gguf_parser.c
index 8b924e7..469122b 100644
--- a/runtime/src/iree/io/formats/gguf/gguf_format.c
+++ b/runtime/src/iree/io/formats/gguf/gguf_parser.c
@@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/io/formats/gguf/gguf_format.h"
+#include "iree/io/formats/gguf/gguf_parser.h"
#include <ctype.h>
@@ -652,7 +652,7 @@
return iree_io_parameter_index_add(parser->index, &entry);
}
-IREE_API_EXPORT iree_status_t iree_io_parse_gguf_index_from_memory(
+static iree_status_t iree_io_parse_gguf_index_from_memory(
iree_io_file_handle_t* file_handle, iree_const_byte_span_t file_contents,
iree_io_parameter_index_t* index) {
// Read the header enough to check for file validity and version.
diff --git a/runtime/src/iree/io/formats/gguf/gguf_format.h b/runtime/src/iree/io/formats/gguf/gguf_parser.h
similarity index 84%
rename from runtime/src/iree/io/formats/gguf/gguf_format.h
rename to runtime/src/iree/io/formats/gguf/gguf_parser.h
index 0e49ad6..fc0064b 100644
--- a/runtime/src/iree/io/formats/gguf/gguf_format.h
+++ b/runtime/src/iree/io/formats/gguf/gguf_parser.h
@@ -4,8 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#ifndef IREE_IO_FORMATS_GGUF_GGUF_FORMAT_H_
-#define IREE_IO_FORMATS_GGUF_GGUF_FORMAT_H_
+#ifndef IREE_IO_FORMATS_GGUF_GGUF_PARSER_H_
+#define IREE_IO_FORMATS_GGUF_GGUF_PARSER_H_
#include "iree/base/api.h"
#include "iree/io/file_handle.h"
@@ -26,4 +26,4 @@
} // extern "C"
#endif // __cplusplus
-#endif // IREE_IO_FORMATS_GGUF_GGUF_FORMAT_H_
+#endif // IREE_IO_FORMATS_GGUF_GGUF_PARSER_H_
diff --git a/runtime/src/iree/io/formats/gguf/gguf_format_test.cc b/runtime/src/iree/io/formats/gguf/gguf_parser_test.cc
similarity index 97%
rename from runtime/src/iree/io/formats/gguf/gguf_format_test.cc
rename to runtime/src/iree/io/formats/gguf/gguf_parser_test.cc
index c02a535..7ecafb2 100644
--- a/runtime/src/iree/io/formats/gguf/gguf_format_test.cc
+++ b/runtime/src/iree/io/formats/gguf/gguf_parser_test.cc
@@ -4,9 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/io/formats/gguf/gguf_format.h"
+#include "iree/io/formats/gguf/gguf_parser.h"
-#include "iree/base/internal/file_io.h"
#include "iree/io/formats/gguf/testdata/gguf_files.h"
#include "iree/testing/gtest.h"
#include "iree/testing/status_matchers.h"
diff --git a/runtime/src/iree/io/formats/irpa/BUILD.bazel b/runtime/src/iree/io/formats/irpa/BUILD.bazel
new file mode 100644
index 0000000..84a1879
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/BUILD.bazel
@@ -0,0 +1,45 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_runtime_cc_library", "iree_runtime_cc_test")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+iree_runtime_cc_library(
+ name = "irpa",
+ srcs = [
+ "irpa_builder.c",
+ "irpa_parser.c",
+ ],
+ hdrs = [
+ "irpa_builder.h",
+ "irpa_parser.h",
+ ],
+ deps = [
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/io:file_handle",
+ "//runtime/src/iree/io:parameter_index",
+ "//runtime/src/iree/io:stream",
+ "//runtime/src/iree/schemas:parameter_archive",
+ ],
+)
+
+iree_runtime_cc_test(
+ name = "irpa_parser_test",
+ srcs = ["irpa_parser_test.cc"],
+ tags = ["requires-filesystem"],
+ deps = [
+ ":irpa",
+ "//runtime/src/iree/base/internal:file_io",
+ "//runtime/src/iree/io/formats/irpa/testdata:irpa_files",
+ "//runtime/src/iree/testing:gtest",
+ "//runtime/src/iree/testing:gtest_main",
+ ],
+)
diff --git a/runtime/src/iree/io/formats/irpa/CMakeLists.txt b/runtime/src/iree/io/formats/irpa/CMakeLists.txt
new file mode 100644
index 0000000..04f7609
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/CMakeLists.txt
@@ -0,0 +1,46 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# runtime/src/iree/io/formats/irpa/BUILD.bazel #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ irpa
+ HDRS
+ "irpa_builder.h"
+ "irpa_parser.h"
+ SRCS
+ "irpa_builder.c"
+ "irpa_parser.c"
+ DEPS
+ iree::base
+ iree::io::file_handle
+ iree::io::parameter_index
+ iree::io::stream
+ iree::schemas::parameter_archive
+ PUBLIC
+)
+
+iree_cc_test(
+ NAME
+ irpa_parser_test
+ SRCS
+ "irpa_parser_test.cc"
+ DEPS
+ ::irpa
+ iree::base::internal::file_io
+ iree::io::formats::irpa::testdata::irpa_files
+ iree::testing::gtest
+ iree::testing::gtest_main
+ LABELS
+ "requires-filesystem"
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/io/formats/irpa/irpa_builder.c b/runtime/src/iree/io/formats/irpa/irpa_builder.c
new file mode 100644
index 0000000..91563e4
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/irpa_builder.c
@@ -0,0 +1,422 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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 "iree/io/formats/irpa/irpa_builder.h"
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_initialize(
+ iree_allocator_t host_allocator,
+ iree_io_parameter_archive_builder_t* out_builder) {
+ IREE_ASSERT_ARGUMENT(out_builder);
+ memset(out_builder, 0, sizeof(*out_builder));
+ out_builder->host_allocator = host_allocator;
+ out_builder->file_alignment =
+ IREE_IO_PARAMETER_ARCHIVE_DEFAULT_FILE_ALIGNMENT;
+ return iree_io_parameter_index_create(host_allocator, &out_builder->index);
+}
+
+IREE_API_EXPORT void iree_io_parameter_archive_builder_deinitialize(
+ iree_io_parameter_archive_builder_t* builder) {
+ IREE_ASSERT_ARGUMENT(builder);
+ iree_io_parameter_index_release(builder->index);
+ memset(builder, 0, sizeof(*builder));
+}
+
+static iree_io_physical_size_t
+iree_io_parameter_archive_builder_storage_alignment(
+ const iree_io_parameter_archive_builder_t* builder) {
+ IREE_ASSERT_ARGUMENT(builder);
+ return builder->storage_alignment ? builder->storage_alignment : 1;
+}
+
+static iree_io_physical_offset_t
+iree_io_parameter_archive_builder_storage_offset(
+ const iree_io_parameter_archive_builder_t* builder) {
+ IREE_ASSERT_ARGUMENT(builder);
+ return iree_align_uint64(
+ iree_align_uint64(sizeof(iree_io_parameter_archive_header_v0_t),
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT) +
+ builder->entry_segment_size + builder->metadata_segment_size,
+ iree_io_parameter_archive_builder_storage_alignment(builder));
+}
+
+IREE_API_EXPORT iree_io_physical_size_t
+iree_io_parameter_archive_builder_total_size(
+ const iree_io_parameter_archive_builder_t* builder) {
+ IREE_ASSERT_ARGUMENT(builder);
+ return iree_align_uint64(
+ iree_io_parameter_archive_builder_storage_offset(builder) +
+ builder->storage_segment_size,
+ builder->file_alignment);
+}
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_write(
+ const iree_io_parameter_archive_builder_t* builder,
+ iree_io_file_handle_t* file_handle, iree_io_physical_offset_t file_offset,
+ iree_io_stream_t* stream, iree_io_parameter_index_t* target_index) {
+ IREE_ASSERT_ARGUMENT(builder);
+ IREE_ASSERT_ARGUMENT(file_handle);
+ IREE_ASSERT_ARGUMENT(stream);
+ IREE_ASSERT_ARGUMENT(target_index);
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ // Calculate relative segment offsets in the file based on the expected header
+ // size. This allows us to emit header-relative file offsets in the header
+ // itself but absolute offsets for consumers of the target_index.
+ const iree_io_parameter_archive_range_t entry_segment = {
+ .offset = iree_align_uint64(sizeof(iree_io_parameter_archive_header_v0_t),
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT),
+ .length = builder->entry_segment_size,
+ };
+ const iree_io_parameter_archive_range_t metadata_segment = {
+ .offset = entry_segment.offset + entry_segment.length,
+ .length = builder->metadata_segment_size,
+ };
+ const iree_io_parameter_archive_range_t storage_segment = {
+ .offset = iree_align_uint64(
+ metadata_segment.offset + metadata_segment.length,
+ iree_io_parameter_archive_builder_storage_alignment(builder)),
+ .length = builder->storage_segment_size,
+ };
+
+ // Write the archive header referencing the other segments in the file.
+ iree_io_parameter_archive_header_v0_t header = {
+ .prefix =
+ {
+ .magic = IREE_IO_PARAMETER_ARCHIVE_MAGIC,
+ .version_major = 0,
+ .version_minor = 0,
+ .header_size = sizeof(header),
+ .next_header_offset = 0,
+ .flags = 0,
+ },
+ .entry_count = iree_io_parameter_index_count(builder->index),
+ .entry_segment = entry_segment,
+ .metadata_segment = metadata_segment,
+ .storage_segment = storage_segment,
+ };
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_write(stream, sizeof(header), &header));
+
+ // Write entry table following the header.
+ // This references ranges in the metadata and storage segment but to preserve
+ // forward-only writes we populate those after writing the table.
+ iree_io_physical_offset_t metadata_offset = 0;
+ for (iree_host_size_t i = 0;
+ i < iree_io_parameter_index_count(builder->index); ++i) {
+ // Align each entry to the base alignment.
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_seek_to_alignment(
+ stream, IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT));
+
+ // Query the source entry template.
+ const iree_io_parameter_index_entry_t* source_entry = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_index_get(builder->index, i, &source_entry));
+
+ // Reserve space in the metadata segment.
+ const iree_io_parameter_archive_metadata_ref_t name_ref = {
+ .offset = metadata_offset,
+ .length = source_entry->key.size,
+ };
+ metadata_offset += name_ref.length;
+ const iree_io_parameter_archive_metadata_ref_t metadata_ref = {
+ .offset = source_entry->metadata.data_length ? metadata_offset : 0,
+ .length = source_entry->metadata.data_length,
+ };
+ metadata_offset += metadata_ref.length;
+
+ // Produce the target archive entry based on the template.
+ iree_io_parameter_index_entry_t target_entry = {
+ .key = source_entry->key,
+ .metadata = source_entry->metadata,
+ .length = source_entry->length,
+ .type = source_entry->type,
+ .storage = source_entry->storage,
+ };
+ switch (source_entry->type) {
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT: {
+ iree_io_parameter_archive_splat_entry_t splat_entry = {
+ .header =
+ {
+ .entry_size = sizeof(splat_entry),
+ .type = IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SPLAT,
+ .flags = 0,
+ .name = name_ref,
+ .metadata = metadata_ref,
+ .minimum_alignment = 0,
+ },
+ .length = target_entry.length,
+ .pattern_length = target_entry.storage.splat.pattern_length,
+ };
+ memcpy(splat_entry.pattern, target_entry.storage.splat.pattern,
+ sizeof(splat_entry.pattern));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_stream_write(stream, sizeof(splat_entry), &splat_entry));
+ break;
+ }
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE: {
+ iree_io_parameter_archive_data_entry_t data_entry = {
+ .header =
+ {
+ .entry_size = sizeof(data_entry),
+ .type = IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_DATA,
+ .flags = 0,
+ .name = name_ref,
+ .metadata = metadata_ref,
+ .minimum_alignment =
+ IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT,
+ },
+ .storage =
+ {
+ .offset = target_entry.storage.file.offset,
+ .length = target_entry.length,
+ },
+ };
+ target_entry.storage.file.handle = file_handle;
+ target_entry.storage.file.offset += storage_segment.offset;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_write(stream, sizeof(data_entry), &data_entry));
+ break;
+ }
+ default: {
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "unhandled entry type %d",
+ (int)source_entry->type));
+ }
+ }
+
+ // Add the entry to the target_index referencing the location in the file
+ // reserved for the entry storage.
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_index_add(target_index, &target_entry));
+ }
+
+ // Write out the metadata table.
+ for (iree_host_size_t i = 0;
+ i < iree_io_parameter_index_count(builder->index); ++i) {
+ // Query the source entry template.
+ const iree_io_parameter_index_entry_t* source_entry = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_index_get(builder->index, i, &source_entry));
+
+ // Write header metadata.
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_write(stream, source_entry->key.size,
+ source_entry->key.data));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_write(stream, source_entry->metadata.data_length,
+ source_entry->metadata.data));
+ }
+
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_add_splat_entry(
+ iree_io_parameter_archive_builder_t* builder, iree_string_view_t name,
+ iree_const_byte_span_t metadata, const void* pattern,
+ uint8_t pattern_length, iree_io_physical_size_t data_length) {
+ IREE_ASSERT_ARGUMENT(builder);
+ IREE_ASSERT_ARGUMENT(pattern);
+ if (pattern_length == 0 ||
+ pattern_length > IREE_IO_PARAMETER_MAX_SPLAT_PATTERN_LENGTH ||
+ !iree_is_power_of_two_uint64(pattern_length)) {
+ return iree_make_status(
+ IREE_STATUS_INVALID_ARGUMENT,
+ "splat pattern length %u invalid; must be 1, 2, 4, 8, or 16 bytes",
+ pattern_length);
+ }
+ IREE_TRACE_ZONE_BEGIN(z0);
+ IREE_TRACE_ZONE_APPEND_TEXT(z0, name.data, name.size);
+ iree_io_parameter_index_entry_t entry = {
+ .key = name,
+ .metadata = metadata,
+ .length = data_length,
+ .type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT,
+ .storage =
+ {
+ .splat =
+ {
+ .pattern = {0}, // set below
+ .pattern_length = pattern_length,
+ },
+ },
+ };
+ memcpy(entry.storage.splat.pattern, pattern, pattern_length);
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_index_add(builder->index, &entry));
+ builder->entry_segment_size =
+ iree_align_uint64(builder->entry_segment_size,
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT) +
+ sizeof(iree_io_parameter_archive_splat_entry_t);
+ builder->metadata_segment_size += name.size + metadata.data_length;
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_add_data_entry(
+ iree_io_parameter_archive_builder_t* builder, iree_string_view_t name,
+ iree_const_byte_span_t metadata, iree_io_physical_size_t minimum_alignment,
+ iree_io_physical_size_t data_length) {
+ IREE_ASSERT_ARGUMENT(builder);
+ IREE_TRACE_ZONE_BEGIN(z0);
+ IREE_TRACE_ZONE_APPEND_TEXT(z0, name.data, name.size);
+ iree_io_parameter_index_entry_t entry = {
+ .key = name,
+ .metadata = metadata,
+ .length = data_length,
+ .type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE,
+ .storage =
+ {
+ .file =
+ {
+ .handle = NULL, // set on commit
+ .offset = iree_align_uint64(builder->storage_segment_size,
+ minimum_alignment),
+ },
+ },
+ };
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_index_add(builder->index, &entry));
+ builder->entry_segment_size =
+ iree_align_uint64(builder->entry_segment_size,
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT) +
+ sizeof(iree_io_parameter_archive_data_entry_t);
+ builder->metadata_segment_size += name.size + metadata.data_length;
+ builder->storage_segment_size = entry.storage.file.offset + entry.length;
+ if (!builder->storage_alignment) {
+ // First entry sets the base alignment.
+ builder->storage_alignment = minimum_alignment;
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+IREE_API_EXPORT iree_status_t iree_io_build_parameter_archive(
+ iree_io_parameter_index_t* source_index,
+ iree_io_parameter_index_t* target_index,
+ iree_io_parameter_archive_file_open_callback_t target_file_open,
+ iree_io_physical_offset_t target_file_offset,
+ iree_allocator_t host_allocator) {
+ IREE_ASSERT_ARGUMENT(source_index);
+ IREE_ASSERT_ARGUMENT(target_index);
+ IREE_ASSERT_ARGUMENT(target_file_open.fn);
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_io_parameter_archive_builder_t builder;
+ iree_io_parameter_archive_builder_initialize(host_allocator, &builder);
+
+ // Declare a parameter for each entry in the index.
+ // This lets us calculate the size we require to store the entry metadata and
+ // its contents (if any). No data is accessed yet.
+ iree_status_t status = iree_ok_status();
+ for (iree_host_size_t i = 0; i < iree_io_parameter_index_count(source_index);
+ ++i) {
+ const iree_io_parameter_index_entry_t* source_entry = NULL;
+ status = iree_io_parameter_index_get(source_index, i, &source_entry);
+ if (!iree_status_is_ok(status)) break;
+ switch (source_entry->type) {
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT:
+ status = iree_io_parameter_archive_builder_add_splat_entry(
+ &builder, source_entry->key, source_entry->metadata,
+ source_entry->storage.splat.pattern,
+ source_entry->storage.splat.pattern_length, source_entry->length);
+ break;
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE:
+ status = iree_io_parameter_archive_builder_add_data_entry(
+ &builder, source_entry->key, source_entry->metadata,
+ IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT,
+ source_entry->length);
+ break;
+ default:
+ status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "unhandled index entry storage type %d",
+ (int)source_entry->type);
+ break;
+ }
+ if (!iree_status_is_ok(status)) break;
+ }
+
+ // Open a file of sufficient size (now that we know it) for writing.
+ iree_io_physical_offset_t archive_offset = iree_align_uint64(
+ target_file_offset, IREE_IO_PARAMETER_ARCHIVE_HEADER_ALIGNMENT);
+ iree_io_physical_size_t archive_length =
+ iree_io_parameter_archive_builder_total_size(&builder);
+ iree_io_file_handle_t* target_file_handle = NULL;
+ if (iree_status_is_ok(status)) {
+ status = target_file_open.fn(target_file_open.user_data, archive_offset,
+ archive_length, &target_file_handle);
+ }
+
+ // Wrap the target file in a stream.
+ iree_io_stream_t* target_stream = NULL;
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_io_stream_open(IREE_IO_STREAM_MODE_WRITABLE, target_file_handle,
+ target_file_offset, host_allocator, &target_stream);
+ }
+
+ // Commit the archive header to the file and produce an index referencing it.
+ // This will allow us to know where to copy file contents.
+ if (iree_status_is_ok(status)) {
+ status = iree_io_parameter_archive_builder_write(
+ &builder, target_file_handle, target_file_offset, target_stream,
+ target_index);
+ }
+
+ // Copy over parameter entry file contents (if any).
+ // This is a slow operation and something we could optimize with lower-level
+ // platform primitives.
+ if (iree_status_is_ok(status)) {
+ for (iree_host_size_t i = 0;
+ i < iree_io_parameter_index_count(source_index); ++i) {
+ const iree_io_parameter_index_entry_t* source_entry = NULL;
+ status = iree_io_parameter_index_get(source_index, i, &source_entry);
+ if (!iree_status_is_ok(status)) break;
+ const iree_io_parameter_index_entry_t* target_entry = NULL;
+ status = iree_io_parameter_index_lookup(target_index, source_entry->key,
+ &target_entry);
+ if (!iree_status_is_ok(status)) break;
+ switch (source_entry->type) {
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT:
+ // No work to do.
+ break;
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE:
+ status = iree_io_stream_seek(
+ target_stream, IREE_IO_STREAM_SEEK_SET,
+ target_file_offset + target_entry->storage.file.offset);
+ if (!iree_status_is_ok(status)) break;
+ status = iree_io_stream_write_file(
+ target_stream, source_entry->storage.file.handle,
+ source_entry->storage.file.offset, target_entry->length,
+ host_allocator);
+ break;
+ default:
+ status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "unhandled index entry storage type %d",
+ (int)source_entry->type);
+ break;
+ }
+ if (!iree_status_is_ok(status)) break;
+ }
+ }
+
+ iree_io_stream_release(target_stream);
+
+ // Flush file contents before returning to the caller (in case they open the
+ // file via a different handle).
+ if (iree_status_is_ok(status)) {
+ status = iree_io_file_handle_flush(target_file_handle);
+ }
+
+ iree_io_file_handle_release(target_file_handle);
+ iree_io_parameter_archive_builder_deinitialize(&builder);
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
diff --git a/runtime/src/iree/io/formats/irpa/irpa_builder.h b/runtime/src/iree/io/formats/irpa/irpa_builder.h
new file mode 100644
index 0000000..22ac538
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/irpa_builder.h
@@ -0,0 +1,124 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+#ifndef IREE_IO_FORMATS_IRPA_IRPA_BUILDER_H_
+#define IREE_IO_FORMATS_IRPA_IRPA_BUILDER_H_
+
+#include "iree/base/api.h"
+#include "iree/io/file_handle.h"
+#include "iree/io/parameter_index.h"
+#include "iree/io/stream.h"
+#include "iree/schemas/parameter_archive.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Utility to build IRPA (IREE Parameter Archive) headers in-memory.
+//
+// Example:
+// iree_io_parameter_archive_builder_t builder;
+// iree_io_parameter_archive_builder_initialize(host_allocator, &builder);
+// iree_io_parameter_archive_builder_add_data_entry(&builder, ...);
+// iree_io_parameter_archive_builder_add_data_entry(&builder, ...);
+// iree_io_parameter_archive_builder_add_data_entry(&builder, ...);
+// total_size = iree_io_parameter_archive_builder_total_size(&builder);
+// << create file of total_size, map into memory >>
+// iree_io_parameter_archive_builder_write(&builder, file, &target_index);
+// << file now contains the full archive header >>
+// << target_index now references the ranges in the file >>
+// << write parameter contents, or don't if leaving uninitialized >>
+// iree_io_parameter_archive_builder_deinitialize(&builder);
+typedef struct iree_io_parameter_archive_builder_t {
+ iree_allocator_t host_allocator;
+ iree_io_parameter_index_t* index;
+ iree_io_physical_size_t file_alignment;
+ iree_io_physical_size_t entry_segment_size;
+ iree_io_physical_size_t metadata_segment_size;
+ iree_io_physical_size_t storage_segment_size;
+ iree_io_physical_size_t storage_alignment;
+} iree_io_parameter_archive_builder_t;
+
+// Initializes a new parameter builder in |out_builder| for use.
+// iree_io_parameter_archive_builder_deinitialize must be called to drop any
+// resources retained while building.
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_initialize(
+ iree_allocator_t host_allocator,
+ iree_io_parameter_archive_builder_t* out_builder);
+
+// Deinitializes |builder| and drops all resources.
+IREE_API_EXPORT void iree_io_parameter_archive_builder_deinitialize(
+ iree_io_parameter_archive_builder_t* builder);
+
+// Returns the total file size required to store the parameter archive header
+// and contents of all added parameters. Adding new parameters will invalidate
+// this value.
+IREE_API_EXPORT iree_io_physical_size_t
+iree_io_parameter_archive_builder_total_size(
+ const iree_io_parameter_archive_builder_t* builder);
+
+// Writes the parameter archive to the given |file_handle|. The file must have
+// at least enough storage to fit iree_io_parameter_archive_builder_total_size.
+// The archive will be written starting at the given |file_offset|.
+// If an optional |target_index| is provided entries for all parameters will be
+// appended to the index referencing the given |file_handle|.
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_write(
+ const iree_io_parameter_archive_builder_t* builder,
+ iree_io_file_handle_t* file_handle, iree_io_physical_offset_t file_offset,
+ iree_io_stream_t* stream, iree_io_parameter_index_t* target_index);
+
+// Adds a new splat entry to |builder|.
+// Splat entries have no physical storage and exist only in the header.
+// |pattern| and |metadata| (if provided) are copied prior to returning.
+// |pattern_length| must be <= 16 (enough for complex<f64>).
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_add_splat_entry(
+ iree_io_parameter_archive_builder_t* builder, iree_string_view_t name,
+ iree_const_byte_span_t metadata, const void* pattern,
+ uint8_t pattern_length, iree_io_physical_size_t data_length);
+
+// Adds a new data entry to |builder|.
+// |metadata| (if provided) is copied prior to returning.
+// Physical storage will be allocated for |data_length| and it will be aligned
+// to at least |minimum_alignment|.
+IREE_API_EXPORT iree_status_t iree_io_parameter_archive_builder_add_data_entry(
+ iree_io_parameter_archive_builder_t* builder, iree_string_view_t name,
+ iree_const_byte_span_t metadata, iree_io_physical_size_t minimum_alignment,
+ iree_io_physical_size_t data_length);
+
+// Callback for opening a file for writing.
+// Implementations need to ensure that at least |archive_length| bytes are
+// available in the file starting at |archive_offset|.
+typedef iree_status_t(IREE_API_PTR* iree_io_parameter_archive_file_open_fn_t)(
+ void* user_data, iree_io_physical_offset_t archive_offset,
+ iree_io_physical_size_t archive_length,
+ iree_io_file_handle_t** out_file_handle);
+
+// A callback issued to open a file.
+typedef struct {
+ // Callback function pointer.
+ iree_io_parameter_archive_file_open_fn_t fn;
+ // User data passed to the callback function. Unowned.
+ void* user_data;
+} iree_io_parameter_archive_file_open_callback_t;
+
+// Builds a parameter archive from the given |source_index| and returns a new
+// index in |target_index| referencing the new archive file.
+// The total size of the archive will be calculated and the provided
+// |target_file_open| callback will be used to acquire a handle to a writeable
+// file with enough capacity to fit the whole archive. All parameter contents
+// will be written and flushed to the file prior to returning.
+IREE_API_EXPORT iree_status_t iree_io_build_parameter_archive(
+ iree_io_parameter_index_t* source_index,
+ iree_io_parameter_index_t* target_index,
+ iree_io_parameter_archive_file_open_callback_t target_file_open,
+ iree_io_physical_offset_t target_file_offset,
+ iree_allocator_t host_allocator);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
+
+#endif // IREE_IO_FORMATS_IRPA_IRPA_BUILDER_H_
diff --git a/runtime/src/iree/io/formats/irpa/irpa_parser.c b/runtime/src/iree/io/formats/irpa/irpa_parser.c
new file mode 100644
index 0000000..04fac9c
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/irpa_parser.c
@@ -0,0 +1,348 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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 "iree/io/formats/irpa/irpa_parser.h"
+
+#include "iree/schemas/parameter_archive.h"
+
+static iree_status_t iree_io_verify_irpa_v0_file_range(
+ iree_const_byte_span_t file_contents, iree_io_physical_offset_t base_offset,
+ iree_io_parameter_archive_range_t range) {
+ if (range.length == 0) return iree_ok_status();
+ if (base_offset + range.offset + range.length > file_contents.data_length) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "file segment out of range (%" PRIu64 " to %" PRIu64
+ " for %" PRIu64 ", file_size=%" PRIhsz ")",
+ base_offset + range.offset,
+ base_offset + range.offset + range.length - 1,
+ range.length, file_contents.data_length);
+ }
+ return iree_ok_status();
+}
+
+static iree_status_t iree_io_resolve_irpa_v0_string(
+ iree_const_byte_span_t file_contents, iree_io_physical_offset_t base_offset,
+ const iree_io_parameter_archive_header_v0_t* header,
+ iree_io_parameter_archive_metadata_ref_t range,
+ iree_string_view_t* out_view) {
+ *out_view = iree_string_view_empty();
+ if (range.length == 0) return iree_ok_status();
+ if (range.offset + range.length > header->metadata_segment.length) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "metadata segment reference out of range (%" PRIu64
+ " to %" PRIu64 " for %" PRIu64
+ ", segment_size=%" PRIu64 ")",
+ range.offset, range.offset + range.length - 1,
+ range.length, header->metadata_segment.length);
+ }
+ iree_io_physical_offset_t view_offset =
+ base_offset + header->metadata_segment.offset + range.offset;
+ *out_view =
+ iree_make_string_view(file_contents.data + view_offset, range.length);
+ return iree_ok_status();
+}
+
+static iree_status_t iree_io_resolve_irpa_v0_metadata(
+ iree_const_byte_span_t file_contents, iree_io_physical_offset_t base_offset,
+ const iree_io_parameter_archive_header_v0_t* header,
+ iree_io_parameter_archive_metadata_ref_t range,
+ iree_const_byte_span_t* out_span) {
+ *out_span = iree_const_byte_span_empty();
+ if (range.length == 0) return iree_ok_status();
+ if (range.offset + range.length > header->metadata_segment.length) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "metadata segment reference out of range (%" PRIu64
+ " to %" PRIu64 " for %" PRIu64
+ ", segment_size=%" PRIu64 ")",
+ range.offset, range.offset + range.length - 1,
+ range.length, header->metadata_segment.length);
+ }
+ iree_io_physical_offset_t span_offset =
+ base_offset + header->metadata_segment.offset + range.offset;
+ *out_span =
+ iree_make_const_byte_span(file_contents.data + span_offset, range.length);
+ return iree_ok_status();
+}
+
+static iree_status_t iree_io_resolve_irpa_v0_storage(
+ iree_const_byte_span_t file_contents, iree_io_physical_offset_t base_offset,
+ const iree_io_parameter_archive_header_v0_t* header,
+ iree_io_parameter_archive_storage_ref_t range,
+ iree_io_physical_offset_t* out_offset) {
+ *out_offset = 0;
+ if (range.length == 0) return iree_ok_status();
+ if (range.offset + range.length > header->storage_segment.length) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "storage segment reference out of range (%" PRIu64
+ " to %" PRIu64 " for %" PRIu64
+ ", segment_size=%" PRIu64 ")",
+ range.offset, range.offset + range.length - 1,
+ range.length, header->storage_segment.length);
+ }
+ iree_io_physical_offset_t storage_offset =
+ base_offset + header->storage_segment.offset + range.offset;
+ *out_offset = storage_offset;
+ return iree_ok_status();
+}
+
+static iree_status_t iree_io_parse_irpa_v0_splat_entry(
+ const iree_io_parameter_archive_header_v0_t* header,
+ const iree_io_parameter_archive_splat_entry_t* splat_entry,
+ iree_string_view_t name, iree_const_byte_span_t metadata,
+ iree_io_parameter_index_t* index) {
+ if (splat_entry->header.entry_size < sizeof(*splat_entry)) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "splat entry length underflow");
+ }
+ if (splat_entry->pattern_length > sizeof(splat_entry->pattern)) {
+ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "splat pattern length %u out of bounds %" PRIhsz,
+ splat_entry->pattern_length,
+ sizeof(splat_entry->pattern));
+ }
+ iree_io_parameter_index_entry_t entry = {
+ .key = name,
+ .metadata = metadata,
+ .length = splat_entry->length,
+ .type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT,
+ .storage =
+ {
+ .splat =
+ {
+ .pattern = {0}, // set below
+ .pattern_length = splat_entry->pattern_length,
+ },
+ },
+ };
+ memcpy(entry.storage.splat.pattern, splat_entry->pattern,
+ entry.storage.splat.pattern_length);
+ return iree_io_parameter_index_add(index, &entry);
+}
+
+static iree_status_t iree_io_parse_irpa_v0_data_entry(
+ iree_io_file_handle_t* file_handle, iree_const_byte_span_t file_contents,
+ iree_io_physical_offset_t base_offset,
+ const iree_io_parameter_archive_header_v0_t* header,
+ const iree_io_parameter_archive_data_entry_t* data_entry,
+ iree_string_view_t name, iree_const_byte_span_t metadata,
+ iree_io_parameter_index_t* index) {
+ if (data_entry->header.entry_size < sizeof(*data_entry)) {
+ return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "data entry length underflow");
+ }
+ iree_io_physical_offset_t storage_offset = 0;
+ IREE_RETURN_IF_ERROR(
+ iree_io_resolve_irpa_v0_storage(file_contents, base_offset, header,
+ data_entry->storage, &storage_offset));
+ iree_io_parameter_index_entry_t entry = {
+ .key = name,
+ .metadata = metadata,
+ .length = data_entry->storage.length,
+ .type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE,
+ .storage =
+ {
+ .file =
+ {
+ .handle = file_handle,
+ .offset = storage_offset,
+ },
+ },
+ };
+ return iree_io_parameter_index_add(index, &entry);
+}
+
+static iree_status_t iree_io_parse_irpa_v0_index_from_memory(
+ iree_io_file_handle_t* file_handle, iree_const_byte_span_t file_contents,
+ iree_io_physical_offset_t base_offset,
+ const iree_io_parameter_archive_header_prefix_t* header_prefix,
+ iree_io_parameter_index_t* index) {
+ // Get the full header struct.
+ if (header_prefix->version_minor > 0) {
+ return iree_make_status(
+ IREE_STATUS_UNIMPLEMENTED,
+ "IRPA version %u.%u not supported (major supported "
+ "but minor is newer than the runtime trying to parse it)",
+ header_prefix->version_major, header_prefix->version_minor);
+ }
+ if (header_prefix->header_size !=
+ sizeof(iree_io_parameter_archive_header_v0_t)) {
+ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "IRPA v0 header expected to be exactly %" PRIhsz
+ " bytes but was reported as %" PRIu64,
+ sizeof(iree_io_parameter_archive_header_v0_t),
+ header_prefix->header_size);
+ }
+ const iree_io_parameter_archive_header_v0_t* header =
+ (const iree_io_parameter_archive_header_v0_t*)file_contents.data;
+
+ // Verify the base data ranges; this lets all subsequent checks be against the
+ // header instead of needing to know about the view into the file.
+ IREE_RETURN_IF_ERROR(iree_io_verify_irpa_v0_file_range(
+ file_contents, base_offset, header->entry_segment),
+ "verifying entry table");
+ IREE_RETURN_IF_ERROR(
+ iree_io_verify_irpa_v0_file_range(file_contents, base_offset,
+ header->metadata_segment),
+ "verifying metadata segment");
+ IREE_RETURN_IF_ERROR(iree_io_verify_irpa_v0_file_range(
+ file_contents, base_offset, header->storage_segment),
+ "verifying storage segment");
+
+ // Walk the entry table, which has variable-length entries.
+ iree_io_physical_offset_t entry_offset =
+ base_offset + header->entry_segment.offset;
+ iree_io_physical_size_t entry_size_remaining = header->entry_segment.length;
+ for (iree_io_physical_size_t i = 0; i < header->entry_count; ++i) {
+ // Ensure there's enough space in the table for the base entry header.
+ if (entry_size_remaining <
+ sizeof(iree_io_parameter_archive_entry_header_t)) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "entry table truncated; insufficient bytes for base entry header");
+ }
+
+ // Ensure there's enough space for the declared entry size (if any larger).
+ const iree_io_parameter_archive_entry_header_t* entry_header =
+ (const iree_io_parameter_archive_entry_header_t*)(file_contents.data +
+ entry_offset);
+ if (entry_header->entry_size < sizeof(*entry_header) ||
+ entry_size_remaining < entry_header->entry_size) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "entry table truncated; insufficient bytes for sized header");
+ }
+ // TODO(benvanik): make this explicit with iree_io_stream_seek_to_alignment.
+ iree_io_physical_offset_t aligned_entry_size = iree_align_uint64(
+ entry_header->entry_size, IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT);
+ entry_offset += aligned_entry_size;
+ entry_size_remaining -= aligned_entry_size;
+
+ // Resolve entry metadata from the archive metadata segment.
+ iree_string_view_t name = iree_string_view_empty();
+ IREE_RETURN_IF_ERROR(
+ iree_io_resolve_irpa_v0_string(file_contents, base_offset, header,
+ entry_header->name, &name),
+ "resolving entry name");
+ iree_const_byte_span_t metadata = iree_const_byte_span_empty();
+ IREE_RETURN_IF_ERROR(
+ iree_io_resolve_irpa_v0_metadata(file_contents, base_offset, header,
+ entry_header->metadata, &metadata),
+ "resolving entry metadata");
+
+ // Handle each entry type.
+ switch (entry_header->type) {
+ case IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SKIP:
+ break;
+ case IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SPLAT: {
+ IREE_RETURN_IF_ERROR(iree_io_parse_irpa_v0_splat_entry(
+ header,
+ (const iree_io_parameter_archive_splat_entry_t*)entry_header, name,
+ metadata, index));
+ break;
+ }
+ case IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_DATA: {
+ IREE_RETURN_IF_ERROR(iree_io_parse_irpa_v0_data_entry(
+ file_handle, file_contents, base_offset, header,
+ (const iree_io_parameter_archive_data_entry_t*)entry_header, name,
+ metadata, index));
+ break;
+ }
+ default:
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "parser does not support entry type %d",
+ (int)entry_header->type);
+ }
+ }
+
+ return iree_ok_status();
+}
+
+static iree_status_t iree_io_parse_irpa_index_from_memory(
+ iree_io_file_handle_t* file_handle, iree_const_byte_span_t file_contents,
+ iree_io_physical_offset_t base_offset, iree_io_parameter_index_t* index) {
+ // Check the basic header information is something we can process.
+ if (file_contents.data_length <
+ base_offset + sizeof(iree_io_parameter_archive_header_prefix_t)) {
+ return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "not enough bytes for a valid IRPA header; file "
+ "may be empty or truncated");
+ }
+ const iree_io_parameter_archive_header_prefix_t* header_prefix =
+ (const iree_io_parameter_archive_header_prefix_t*)(file_contents.data +
+ base_offset);
+ if (header_prefix->magic != IREE_IO_PARAMETER_ARCHIVE_MAGIC) {
+ return iree_make_status(
+ IREE_STATUS_INVALID_ARGUMENT,
+ "IRPA file magic missing or invalid %08X; expected %08X",
+ header_prefix->magic, IREE_IO_PARAMETER_ARCHIVE_MAGIC);
+ }
+ if (header_prefix->header_size > file_contents.data_length) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "file buffer underrun parsing header of reported size %" PRIu64
+ " (only %" PRIhsz " bytes available)",
+ header_prefix->header_size, file_contents.data_length);
+ }
+ if (header_prefix->next_header_offset != 0 &&
+ file_contents.data_length <
+ base_offset + header_prefix->next_header_offset +
+ sizeof(iree_io_parameter_archive_header_prefix_t)) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "file buffer underrun verifying linked header at offset %" PRIu64
+ " (only %" PRIhsz " bytes available)",
+ base_offset + header_prefix->next_header_offset,
+ file_contents.data_length);
+ }
+
+ // Route major versions to their parsers, allowing us to change everything but
+ // the prefix without breaking compatibility.
+ switch (header_prefix->version_major) {
+ case 0: {
+ IREE_RETURN_IF_ERROR(iree_io_parse_irpa_v0_index_from_memory(
+ file_handle, file_contents, base_offset, header_prefix, index));
+ break;
+ }
+ default: {
+ return iree_make_status(
+ IREE_STATUS_UNIMPLEMENTED,
+ "IRPA major version %u.%u not supported by this runtime",
+ header_prefix->version_major, header_prefix->version_minor);
+ }
+ }
+
+ // If there's a linked header then tail-call process it.
+ if (header_prefix->next_header_offset == 0) return iree_ok_status();
+ return iree_io_parse_irpa_index_from_memory(
+ file_handle, file_contents,
+ base_offset + header_prefix->next_header_offset, index);
+}
+
+IREE_API_EXPORT iree_status_t iree_io_parse_irpa_index(
+ iree_io_file_handle_t* file_handle, iree_io_parameter_index_t* index) {
+ IREE_ASSERT_ARGUMENT(index);
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ // Today we only support memory files.
+ // TODO(benvanik): support iree_io_stream_t wrapping for parsing the index.
+ if (iree_io_file_handle_type(file_handle) !=
+ IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION) {
+ IREE_TRACE_ZONE_END(z0);
+ return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
+ "non-memory irpa files not yet supported");
+ }
+ iree_byte_span_t host_allocation =
+ iree_io_file_handle_primitive(file_handle).value.host_allocation;
+
+ iree_status_t status = iree_io_parse_irpa_index_from_memory(
+ file_handle,
+ iree_make_const_byte_span(host_allocation.data,
+ host_allocation.data_length),
+ /*base_offset=*/0, index);
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
diff --git a/runtime/src/iree/io/formats/irpa/irpa_parser.h b/runtime/src/iree/io/formats/irpa/irpa_parser.h
new file mode 100644
index 0000000..4423105
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/irpa_parser.h
@@ -0,0 +1,26 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+#ifndef IREE_IO_FORMATS_IRPA_irpa_parser_H_
+#define IREE_IO_FORMATS_IRPA_irpa_parser_H_
+
+#include "iree/base/api.h"
+#include "iree/io/file_handle.h"
+#include "iree/io/parameter_index.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Parses an IREE archive file and merges its contained resources into |index|.
+IREE_API_EXPORT iree_status_t iree_io_parse_irpa_index(
+ iree_io_file_handle_t* file_handle, iree_io_parameter_index_t* index);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
+
+#endif // IREE_IO_FORMATS_IRPA_irpa_parser_H_
diff --git a/runtime/src/iree/io/formats/irpa/irpa_parser_test.cc b/runtime/src/iree/io/formats/irpa/irpa_parser_test.cc
new file mode 100644
index 0000000..944c2b4
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/irpa_parser_test.cc
@@ -0,0 +1,157 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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 "iree/io/formats/irpa/irpa_parser.h"
+
+#include "iree/io/formats/irpa/testdata/irpa_files.h"
+#include "iree/testing/gtest.h"
+#include "iree/testing/status_matchers.h"
+
+namespace iree {
+namespace {
+
+static iree_io_file_handle_t* OpenTestFile(const char* name) {
+ const struct iree_file_toc_t* file_toc = iree_io_irpa_files_create();
+ for (size_t i = 0; i < iree_io_irpa_files_size(); ++i) {
+ if (strcmp(file_toc[i].name, name) == 0) {
+ iree_io_file_handle_t* file_handle = NULL;
+ IREE_CHECK_OK(iree_io_file_handle_wrap_host_allocation(
+ IREE_IO_FILE_ACCESS_READ,
+ iree_make_byte_span((void*)file_toc[i].data, file_toc[i].size),
+ iree_io_file_handle_release_callback_null(), iree_allocator_system(),
+ &file_handle));
+ return file_handle;
+ }
+ }
+ IREE_CHECK_OK(iree_make_status(
+ IREE_STATUS_NOT_FOUND,
+ "test file `%s` not found embedded into test binary", name));
+ return NULL;
+}
+
+TEST(IrpaFormatTest, Empty) {
+ iree_io_parameter_index_t* index = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_create(iree_allocator_system(), &index));
+
+ iree_io_file_handle_t* file_handle = OpenTestFile("empty.irpa");
+ IREE_ASSERT_OK(iree_io_parse_irpa_index(file_handle, index));
+ EXPECT_EQ(0, iree_io_parameter_index_count(index));
+ iree_io_file_handle_release(file_handle);
+
+ iree_io_parameter_index_release(index);
+}
+
+TEST(IrpaFormatTest, SingleParameters) {
+ iree_io_parameter_index_t* index = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_create(iree_allocator_system(), &index));
+
+ iree_io_file_handle_t* file_handle = OpenTestFile("single.irpa");
+ IREE_ASSERT_OK(iree_io_parse_irpa_index(file_handle, index));
+ EXPECT_EQ(1, iree_io_parameter_index_count(index));
+ iree_io_file_handle_release(file_handle);
+
+ const iree_io_parameter_index_entry_t* entry0 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key0"), &entry0));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key0"), entry0->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry0->metadata));
+ EXPECT_EQ(entry0->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE);
+ EXPECT_EQ(entry0->storage.file.offset, 192);
+ EXPECT_EQ(entry0->length, 16);
+
+ iree_io_parameter_index_release(index);
+}
+
+TEST(IrpaFormatTest, MultipleParameters) {
+ iree_io_parameter_index_t* index = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_create(iree_allocator_system(), &index));
+
+ iree_io_file_handle_t* file_handle = OpenTestFile("multiple.irpa");
+ IREE_ASSERT_OK(iree_io_parse_irpa_index(file_handle, index));
+ EXPECT_EQ(2, iree_io_parameter_index_count(index));
+ iree_io_file_handle_release(file_handle);
+
+ const iree_io_parameter_index_entry_t* entry0 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key0"), &entry0));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key0"), entry0->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry0->metadata));
+ EXPECT_EQ(entry0->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE);
+ EXPECT_EQ(entry0->storage.file.offset, 320);
+ EXPECT_EQ(entry0->length, 16);
+
+ const iree_io_parameter_index_entry_t* entry1 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key1"), &entry1));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key1"), entry1->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry1->metadata));
+ EXPECT_EQ(entry1->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE);
+ EXPECT_EQ(entry1->storage.file.offset, 384);
+ EXPECT_EQ(entry1->length, 5);
+
+ iree_io_parameter_index_release(index);
+}
+
+TEST(IrpaFormatTest, MixedDataAndSplats) {
+ iree_io_parameter_index_t* index = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_create(iree_allocator_system(), &index));
+
+ iree_io_file_handle_t* file_handle = OpenTestFile("mixed.irpa");
+ IREE_ASSERT_OK(iree_io_parse_irpa_index(file_handle, index));
+ EXPECT_EQ(4, iree_io_parameter_index_count(index));
+ iree_io_file_handle_release(file_handle);
+
+ const iree_io_parameter_index_entry_t* entry0 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key0"), &entry0));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key0"), entry0->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry0->metadata));
+ EXPECT_EQ(entry0->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE);
+ EXPECT_EQ(entry0->storage.file.offset, 512);
+ EXPECT_EQ(entry0->length, 16);
+
+ const iree_io_parameter_index_entry_t* entry1 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key1"), &entry1));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key1"), entry1->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry1->metadata));
+ EXPECT_EQ(entry1->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE);
+ EXPECT_EQ(entry1->storage.file.offset, 576);
+ EXPECT_EQ(entry1->length, 5);
+
+ const iree_io_parameter_index_entry_t* entry2 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key2"), &entry2));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key2"), entry2->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry2->metadata));
+ EXPECT_EQ(entry2->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT);
+ EXPECT_EQ(entry2->storage.splat.pattern_length, 1);
+ const int8_t entry2_pattern = 102;
+ EXPECT_EQ(0, memcmp(&entry2_pattern, entry2->storage.splat.pattern,
+ sizeof(entry2_pattern)));
+ EXPECT_EQ(entry2->length, 1);
+
+ const iree_io_parameter_index_entry_t* entry3 = NULL;
+ IREE_ASSERT_OK(
+ iree_io_parameter_index_lookup(index, IREE_SV("key3"), &entry3));
+ EXPECT_TRUE(iree_string_view_equal(IREE_SV("key3"), entry3->key));
+ EXPECT_TRUE(iree_const_byte_span_is_empty(entry3->metadata));
+ EXPECT_EQ(entry3->type, IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT);
+ EXPECT_EQ(entry3->storage.splat.pattern_length, 8);
+ const int64_t entry3_pattern = 9223372036854775807ll;
+ EXPECT_EQ(0, memcmp(&entry3_pattern, entry3->storage.splat.pattern,
+ sizeof(entry3_pattern)));
+ EXPECT_EQ(entry3->length, 33554432);
+
+ iree_io_parameter_index_release(index);
+}
+
+} // namespace
+} // namespace iree
diff --git a/runtime/src/iree/io/formats/irpa/testdata/BUILD.bazel b/runtime/src/iree/io/formats/irpa/testdata/BUILD.bazel
new file mode 100644
index 0000000..2a5f86e
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/BUILD.bazel
@@ -0,0 +1,23 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+load("//build_tools/embed_data:build_defs.bzl", "c_embed_data")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+c_embed_data(
+ name = "irpa_files",
+ testonly = True,
+ srcs = glob(["*.irpa"]),
+ c_file_output = "irpa_files.c",
+ flatten = True,
+ h_file_output = "irpa_files.h",
+ identifier = "iree_io_irpa_files",
+)
diff --git a/runtime/src/iree/io/formats/irpa/testdata/CMakeLists.txt b/runtime/src/iree/io/formats/irpa/testdata/CMakeLists.txt
new file mode 100644
index 0000000..0e3662d
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/CMakeLists.txt
@@ -0,0 +1,30 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# runtime/src/iree/io/formats/irpa/testdata/BUILD.bazel #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+file(GLOB _GLOB_X_IRPA LIST_DIRECTORIES false RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} CONFIGURE_DEPENDS *.irpa)
+iree_c_embed_data(
+ NAME
+ irpa_files
+ SRCS
+ "${_GLOB_X_IRPA}"
+ C_FILE_OUTPUT
+ "irpa_files.c"
+ H_FILE_OUTPUT
+ "irpa_files.h"
+ IDENTIFIER
+ "iree_io_irpa_files"
+ TESTONLY
+ FLATTEN
+ PUBLIC
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/io/formats/irpa/testdata/empty.irpa b/runtime/src/iree/io/formats/irpa/testdata/empty.irpa
new file mode 100644
index 0000000..2c8a084
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/empty.irpa
Binary files differ
diff --git a/runtime/src/iree/io/formats/irpa/testdata/generate_irpa_files.sh b/runtime/src/iree/io/formats/irpa/testdata/generate_irpa_files.sh
new file mode 100644
index 0000000..9c28ab9
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/generate_irpa_files.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+# Copyright 2023 The IREE Authors
+#
+# Licensed 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
+
+# TODO(benvanik): replace with python utils using numpy so we can get more than
+# just splat data tensors. For now there's no python writer.
+
+# To regenerate, build iree-create-parameters and run from the project root:
+# $ ./runtime/src/iree/io/formats/irpa/testdata/generate_irpa_files.sh
+#
+# If iree-create-parameters is not on your path you can set the env var:
+# $ IREE_CREATE_PARAMETERS=../iree-create-parameters \
+# ./runtime/src/iree/io/formats/irpa/testdata/generate_irpa_files.sh
+
+# Uncomment to see the iree-create-parameters commands issued:
+# set -x
+set -e
+
+ROOT_DIR=$(git rev-parse --show-toplevel)
+TESTDATA="${ROOT_DIR}/runtime/src/iree/io/formats/irpa/testdata"
+CREATE_PARAMETERS="${IREE_CREATE_PARAMETERS:-iree-create-parameters}"
+
+CMD=(
+ ${CREATE_PARAMETERS}
+ --output=${TESTDATA}/empty.irpa
+)
+"${CMD[@]}"
+
+CMD=(
+ ${CREATE_PARAMETERS}
+ --output=${TESTDATA}/single.irpa
+ --data=key0=4xf32=100.1
+)
+"${CMD[@]}"
+
+CMD=(
+ ${CREATE_PARAMETERS}
+ --output=${TESTDATA}/multiple.irpa
+ --data=key0=4xf32=100.1
+ --data=key1=5xi8=101
+)
+"${CMD[@]}"
+
+CMD=(
+ ${CREATE_PARAMETERS}
+ --output=${TESTDATA}/mixed.irpa
+ --data=key0=4xf32=100.1
+ --data=key1=5xi8=101
+ --splat=key2=i8=102
+ --splat=key3=4096x1024xi64=103
+)
+"${CMD[@]}"
diff --git a/runtime/src/iree/io/formats/irpa/testdata/mixed.irpa b/runtime/src/iree/io/formats/irpa/testdata/mixed.irpa
new file mode 100644
index 0000000..09da720
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/mixed.irpa
Binary files differ
diff --git a/runtime/src/iree/io/formats/irpa/testdata/multiple.irpa b/runtime/src/iree/io/formats/irpa/testdata/multiple.irpa
new file mode 100644
index 0000000..b0853e2
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/multiple.irpa
Binary files differ
diff --git a/runtime/src/iree/io/formats/irpa/testdata/single.irpa b/runtime/src/iree/io/formats/irpa/testdata/single.irpa
new file mode 100644
index 0000000..4c9407e
--- /dev/null
+++ b/runtime/src/iree/io/formats/irpa/testdata/single.irpa
Binary files differ
diff --git a/runtime/src/iree/io/formats/safetensors/BUILD.bazel b/runtime/src/iree/io/formats/safetensors/BUILD.bazel
index a0d6a5e..175ecc4 100644
--- a/runtime/src/iree/io/formats/safetensors/BUILD.bazel
+++ b/runtime/src/iree/io/formats/safetensors/BUILD.bazel
@@ -15,10 +15,10 @@
iree_runtime_cc_library(
name = "safetensors",
srcs = [
- "safetensors_format.c",
+ "safetensors_parser.c",
],
hdrs = [
- "safetensors_format.h",
+ "safetensors_parser.h",
],
deps = [
"//runtime/src/iree/base",
@@ -28,12 +28,10 @@
)
iree_runtime_cc_test(
- name = "safetensors_format_test",
- srcs = ["safetensors_format_test.cc"],
- tags = ["requires-filesystem"],
+ name = "safetensors_parser_test",
+ srcs = ["safetensors_parser_test.cc"],
deps = [
":safetensors",
- "//runtime/src/iree/base/internal:file_io",
"//runtime/src/iree/io/formats/safetensors/testdata:safetensors_files",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
diff --git a/runtime/src/iree/io/formats/safetensors/CMakeLists.txt b/runtime/src/iree/io/formats/safetensors/CMakeLists.txt
index f52ac03..2f82057 100644
--- a/runtime/src/iree/io/formats/safetensors/CMakeLists.txt
+++ b/runtime/src/iree/io/formats/safetensors/CMakeLists.txt
@@ -14,9 +14,9 @@
NAME
safetensors
HDRS
- "safetensors_format.h"
+ "safetensors_parser.h"
SRCS
- "safetensors_format.c"
+ "safetensors_parser.c"
DEPS
iree::base
iree::io::file_handle
@@ -26,17 +26,14 @@
iree_cc_test(
NAME
- safetensors_format_test
+ safetensors_parser_test
SRCS
- "safetensors_format_test.cc"
+ "safetensors_parser_test.cc"
DEPS
::safetensors
- iree::base::internal::file_io
iree::io::formats::safetensors::testdata::safetensors_files
iree::testing::gtest
iree::testing::gtest_main
- LABELS
- "requires-filesystem"
)
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/io/formats/safetensors/safetensors_format.c b/runtime/src/iree/io/formats/safetensors/safetensors_parser.c
similarity index 98%
rename from runtime/src/iree/io/formats/safetensors/safetensors_format.c
rename to runtime/src/iree/io/formats/safetensors/safetensors_parser.c
index a2d215d..04eda05 100644
--- a/runtime/src/iree/io/formats/safetensors/safetensors_format.c
+++ b/runtime/src/iree/io/formats/safetensors/safetensors_parser.c
@@ -4,7 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/io/formats/safetensors/safetensors_format.h"
+#include "iree/io/formats/safetensors/safetensors_parser.h"
#include <ctype.h>
@@ -455,7 +455,7 @@
return iree_io_parameter_index_add(entry_state->index, &entry);
}
-IREE_API_EXPORT iree_status_t iree_io_parse_safetensors_index_from_memory(
+static iree_status_t iree_io_parse_safetensors_index_from_memory(
iree_io_file_handle_t* file_handle, iree_const_byte_span_t file_contents,
iree_io_parameter_index_t* index) {
// Reads the header JSON blob out of the file contents and calculates the base
diff --git a/runtime/src/iree/io/formats/safetensors/safetensors_format.h b/runtime/src/iree/io/formats/safetensors/safetensors_parser.h
similarity index 89%
rename from runtime/src/iree/io/formats/safetensors/safetensors_format.h
rename to runtime/src/iree/io/formats/safetensors/safetensors_parser.h
index 30a6638..c32c104 100644
--- a/runtime/src/iree/io/formats/safetensors/safetensors_format.h
+++ b/runtime/src/iree/io/formats/safetensors/safetensors_parser.h
@@ -4,8 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#ifndef IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_FORMAT_H_
-#define IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_FORMAT_H_
+#ifndef IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_PARSER_H_
+#define IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_PARSER_H_
#include "iree/base/api.h"
#include "iree/io/file_handle.h"
@@ -38,4 +38,4 @@
} // extern "C"
#endif // __cplusplus
-#endif // IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_FORMAT_H_
+#endif // IREE_IO_FORMATS_SAFETENSORS_SAFETENSORS_PARSER_H_
diff --git a/runtime/src/iree/io/formats/safetensors/safetensors_format_test.cc b/runtime/src/iree/io/formats/safetensors/safetensors_parser_test.cc
similarity index 97%
rename from runtime/src/iree/io/formats/safetensors/safetensors_format_test.cc
rename to runtime/src/iree/io/formats/safetensors/safetensors_parser_test.cc
index b1620a4..44caf7b 100644
--- a/runtime/src/iree/io/formats/safetensors/safetensors_format_test.cc
+++ b/runtime/src/iree/io/formats/safetensors/safetensors_parser_test.cc
@@ -4,9 +4,8 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-#include "iree/io/formats/safetensors/safetensors_format.h"
+#include "iree/io/formats/safetensors/safetensors_parser.h"
-#include "iree/base/internal/file_io.h"
#include "iree/io/formats/safetensors/testdata/safetensors_files.h"
#include "iree/testing/gtest.h"
#include "iree/testing/status_matchers.h"
diff --git a/runtime/src/iree/io/parameter_index.c b/runtime/src/iree/io/parameter_index.c
index b1a3e72..13228a7 100644
--- a/runtime/src/iree/io/parameter_index.c
+++ b/runtime/src/iree/io/parameter_index.c
@@ -243,3 +243,86 @@
IREE_TRACE_ZONE_END(z0);
return status;
}
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_index_dump(
+ iree_string_view_t scope, iree_io_parameter_index_t* index,
+ iree_string_builder_t* builder) {
+ iree_host_size_t entry_count = iree_io_parameter_index_count(index);
+ uint64_t total_bytes = 0;
+ for (iree_host_size_t i = 0; i < entry_count; ++i) {
+ const iree_io_parameter_index_entry_t* entry = NULL;
+ IREE_RETURN_IF_ERROR(iree_io_parameter_index_get(index, i, &entry));
+ total_bytes += entry->length;
+ }
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(
+ builder,
+ "//"
+ "===-----------------------------------------------------------------"
+ "---------------------------------------------===//\n"));
+ IREE_RETURN_IF_ERROR(
+ iree_string_builder_append_cstring(builder, "// Parameter scope "));
+ if (iree_string_view_is_empty(scope)) {
+ IREE_RETURN_IF_ERROR(
+ iree_string_builder_append_cstring(builder, "<global>"));
+ } else {
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
+ builder, "`%.*s`", (int)scope.size, scope.data));
+ }
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
+ builder, " (%" PRIhsz " entries, %" PRIu64 " total bytes)\n", entry_count,
+ total_bytes));
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(
+ builder,
+ "//===------------+------------------+------------------+------------"
+ "-----------------------------------------------===//\n"));
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(
+ builder,
+ "// Start | End | Length | Key\n"));
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(
+ builder,
+ "//---------------+------------------+------------------+------------"
+ "--------------------------------------------------//\n"));
+ for (iree_host_size_t i = 0; i < entry_count; ++i) {
+ const iree_io_parameter_index_entry_t* entry = NULL;
+ IREE_RETURN_IF_ERROR(iree_io_parameter_index_get(index, i, &entry));
+ switch (entry->type) {
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT: {
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
+ builder,
+ " - | - | %16" PRIu64 " | `%.*s`\n",
+ entry->length, (int)entry->key.size, entry->key.data));
+ break;
+ }
+ case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE: {
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
+ builder, "%16" PRIu64 " | %16" PRIu64 " | %16" PRIu64 " | `%.*s`\n",
+ entry->storage.file.offset,
+ entry->storage.file.offset + entry->length, entry->length,
+ (int)entry->key.size, entry->key.data));
+ break;
+ }
+ default: {
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_format(
+ builder,
+ " ? | ? | %16" PRIu64 " | `%.*s`\n",
+ entry->length, (int)entry->key.size, entry->key.data));
+ break;
+ }
+ }
+ }
+ IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "\n"));
+ return iree_ok_status();
+}
+
+IREE_API_EXPORT iree_status_t iree_io_parameter_index_fprint(
+ FILE* file, iree_string_view_t scope, iree_io_parameter_index_t* index) {
+ iree_string_builder_t builder;
+ iree_string_builder_initialize(index->host_allocator, &builder);
+ iree_status_t status = iree_io_parameter_index_dump(scope, index, &builder);
+ if (iree_status_is_ok(status)) {
+ fprintf(file, "%.*s", (int)iree_string_builder_size(&builder),
+ iree_string_builder_buffer(&builder));
+ }
+ iree_string_builder_deinitialize(&builder);
+ return status;
+}
diff --git a/runtime/src/iree/io/parameter_index.h b/runtime/src/iree/io/parameter_index.h
index 96ddf66..8f25bf1 100644
--- a/runtime/src/iree/io/parameter_index.h
+++ b/runtime/src/iree/io/parameter_index.h
@@ -25,6 +25,11 @@
IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE,
} iree_io_parameter_index_entry_storage_type_t;
+// Power of two; enough bytes to fit complex128 (complex<f64>).
+// Prefer 1, 2, and 4 byte patterns as they can often hit hardware accelerated
+// fast paths while 8 and 16 may require emulation.
+#define IREE_IO_PARAMETER_MAX_SPLAT_PATTERN_LENGTH 16
+
// An entry in an in-memory file index.
typedef struct iree_io_parameter_index_entry_t {
// Key used to reference this file.
@@ -38,7 +43,7 @@
// Defines the backing storage of a parameter based on its type.
union {
// Describes a synthetic parameter comprised of some repeated sequence of
- // bytes. Sized to allow for up to complex128 (2xf64) values.
+ // bytes. Sized to allow for up to complex128 (complex<f64>) values.
// Valid when type is IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT.
struct {
// Length of the pattern in bytes defining how many bytes of the pattern
@@ -49,7 +54,7 @@
// 0xAA: pattern_length=1 pattern=[0xAA ...]
// 0xBBAA: pattern_length=2 pattern=[0xAA 0xBB ...]
// 0xDDCCBBAA: pattern_length=4 pattern=[0xAA 0xBB 0xCC 0xDD ...]
- uint8_t pattern[16];
+ uint8_t pattern[IREE_IO_PARAMETER_MAX_SPLAT_PATTERN_LENGTH];
} splat;
// Describes a file-backed parameter.
// Valid when type is IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE.
@@ -115,6 +120,17 @@
iree_io_parameter_index_t* index, iree_string_view_t key,
const iree_io_parameter_index_entry_t** out_entry);
+// Formats a textual dump of the parameter |index| to |builder|.
+// An optional |scope| name can be provided to include in the dump.
+IREE_API_EXPORT iree_status_t iree_io_parameter_index_dump(
+ iree_string_view_t scope, iree_io_parameter_index_t* index,
+ iree_string_builder_t* builder);
+
+// Writes a textual dump of the parameter |index| to |file|.
+// An optional |scope| name can be provided to include in the dump.
+IREE_API_EXPORT iree_status_t iree_io_parameter_index_fprint(
+ FILE* file, iree_string_view_t scope, iree_io_parameter_index_t* index);
+
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
diff --git a/runtime/src/iree/io/parameter_provider.h b/runtime/src/iree/io/parameter_provider.h
index e398037..76ce5da 100644
--- a/runtime/src/iree/io/parameter_provider.h
+++ b/runtime/src/iree/io/parameter_provider.h
@@ -76,8 +76,7 @@
// Loads a parameter from |provider| for use on |device|.
// |source_scope| and |source_key| define the parameter and |target_params|
-// defines how the buffer is to be allocated. If the parameter is smaller than
-// |length| any remaining space must be zero-filled.
+// defines how the buffer is to be allocated.
//
// If the implementation is able to meet the expected |target_params| with an
// existing buffer it may be returned without a new allocation. If access allows
@@ -100,8 +99,7 @@
// Reads a parameter from |provider| for use on |device|.
// |source_scope| and |source_key| define the parameter to be read into
-// |target_buffer| at |target_offset|. If the parameter is smaller than
-// |length| any remaining space must be zero-filled.
+// |target_buffer| at |target_offset|.
//
// Returns IREE_STATUS_NOT_FOUND if the parameter is not found.
IREE_API_EXPORT iree_status_t iree_io_parameter_provider_read(
@@ -140,10 +138,9 @@
// Gathers zero or more spans from |provider| into the given |target_buffer|.
// The |enumerator| defines the source keys in |source_scope| and the offset and
-// length in the |target_buffer| of each span. For any parameter is smaller than
-// the length specified by the span any remaining space must be zero-filled.
-// Multiple spans may reference the same source parameter but behavior is
-// undefined if multiple span target ranges overlap.
+// length in the |target_buffer| of each span. Multiple spans may reference the
+// same source parameter but behavior is undefined if multiple span target
+// ranges overlap.
//
// Returns IREE_STATUS_NOT_FOUND if any parameter is not found.
IREE_API_EXPORT iree_status_t iree_io_parameter_provider_gather(
diff --git a/runtime/src/iree/io/scope_map.c b/runtime/src/iree/io/scope_map.c
index 806397b..6b22455 100644
--- a/runtime/src/iree/io/scope_map.c
+++ b/runtime/src/iree/io/scope_map.c
@@ -75,3 +75,13 @@
IREE_TRACE_ZONE_END(z0);
return status;
}
+
+IREE_API_EXPORT iree_status_t iree_io_scope_map_dump(
+ iree_io_scope_map_t* scope_map, iree_string_builder_t* builder) {
+ for (iree_host_size_t i = 0; i < scope_map->count; ++i) {
+ iree_string_view_t scope = scope_map->entries[i]->scope;
+ iree_io_parameter_index_t* index = scope_map->entries[i]->index;
+ IREE_RETURN_IF_ERROR(iree_io_parameter_index_dump(scope, index, builder));
+ }
+ return iree_ok_status();
+}
diff --git a/runtime/src/iree/io/scope_map.h b/runtime/src/iree/io/scope_map.h
index 30e7377..5c46976 100644
--- a/runtime/src/iree/io/scope_map.h
+++ b/runtime/src/iree/io/scope_map.h
@@ -38,6 +38,9 @@
iree_io_scope_map_t* scope_map, iree_string_view_t scope,
iree_io_parameter_index_t** out_index);
+IREE_API_EXPORT iree_status_t iree_io_scope_map_dump(
+ iree_io_scope_map_t* scope_map, iree_string_builder_t* builder);
+
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
diff --git a/runtime/src/iree/io/stream.c b/runtime/src/iree/io/stream.c
index 1fd29de..08471a7 100644
--- a/runtime/src/iree/io/stream.c
+++ b/runtime/src/iree/io/stream.c
@@ -8,7 +8,7 @@
// This is arbitrary - we should allow dynamic block sizes and such but keeping
// this small only requires a reasonable fixed size stack alloc.
-#define IREE_IO_STREAM_COPY_BLOCK_SIZE (16 * 1024)
+#define IREE_IO_STREAM_COPY_BLOCK_SIZE (32 * 1024)
//===----------------------------------------------------------------------===//
// Stream utilities
diff --git a/runtime/src/iree/schemas/BUILD.bazel b/runtime/src/iree/schemas/BUILD.bazel
index 9133efc..294c793 100644
--- a/runtime/src/iree/schemas/BUILD.bazel
+++ b/runtime/src/iree/schemas/BUILD.bazel
@@ -81,3 +81,8 @@
name = "cpu_data",
hdrs = cpu_data_headers,
)
+
+iree_runtime_cc_library(
+ name = "parameter_archive",
+ hdrs = ["parameter_archive.h"],
+)
diff --git a/runtime/src/iree/schemas/CMakeLists.txt b/runtime/src/iree/schemas/CMakeLists.txt
index 4e959fd..776616e 100644
--- a/runtime/src/iree/schemas/CMakeLists.txt
+++ b/runtime/src/iree/schemas/CMakeLists.txt
@@ -110,4 +110,14 @@
PUBLIC
)
+iree_cc_library(
+ NAME
+ parameter_archive
+ HDRS
+ "parameter_archive.h"
+ DEPS
+
+ PUBLIC
+)
+
### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/runtime/src/iree/schemas/parameter_archive.h b/runtime/src/iree/schemas/parameter_archive.h
new file mode 100644
index 0000000..5a67ddd
--- /dev/null
+++ b/runtime/src/iree/schemas/parameter_archive.h
@@ -0,0 +1,241 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+#ifndef IREE_SCHEMAS_PARAMETER_ARCHIVE_H_
+#define IREE_SCHEMAS_PARAMETER_ARCHIVE_H_
+
+#include <stdint.h>
+
+//===----------------------------------------------------------------------===//
+// IREE Parameter Archive
+//===----------------------------------------------------------------------===//
+//
+// Each parameter archive stores zero or more parameters that are referenced by
+// a string name. A single file may contain multiple archives linked together.
+// Archives are designed for optimized concatenation (extend a file and link
+// each archive header), appending (link new archive header), erasing (change
+// entry to skip), and replacing (existing skip and then a new archive header).
+//
+// During mutation the archives may get spread out over the file and hurt
+// startup performance. Repacking archive files with `iree-convert-parameters`
+// will combine all archive headers at the head of the file to avoid needing to
+// scan large files to index their contents.
+//
+// To enable concatenation all offset fields in an archive header are relative
+// to the offset of the header added to the header-relative offset to the
+// appropriate segment. Strings, paths, and metadata blobs are stored in a
+// metadata segment that has no alignment specified. Parameters contents are
+// stored separately in a data storage segment that meets the alignment
+// requirements specified by each parameter.
+//
+// For aiding testing/benchmarking workflows involving large parameters the
+// archive can contain entries backed by a splatted value instead of any real
+// data. This allows for a program that references parameters to be compiled and
+// run against a "fake" set of splatted parameters instead of needing the full
+// original file on each machine running the tests/benchmarks.
+// Use `iree-convert-parameters` with the `--strip` flag to strip all parameter
+// values or `--splat=key` to strip selected parameters.
+
+#if defined(_MSC_VER)
+#define IREE_IO_PACKED_BEGIN __pragma(pack(push, 1))
+#define IREE_IO_PACKED_END __pragma(pack(pop))
+#else
+#define IREE_IO_PACKED_BEGIN _Pragma("pack(push, 1)")
+#define IREE_IO_PACKED_END _Pragma("pack(pop)")
+#endif // _MSC_VER
+
+IREE_IO_PACKED_BEGIN
+
+typedef uint64_t iree_io_physical_offset_t;
+typedef uint64_t iree_io_physical_size_t;
+
+// Preferred alignment for the file size containing an archive.
+// This is the (common) small page size on most platforms and required in order
+// to ensure fast-path memory mapping. Large pages and other configuration may
+// require a larger alignment but we leave that to those special cases to
+// handle.
+#define IREE_IO_PARAMETER_ARCHIVE_DEFAULT_FILE_ALIGNMENT 4096
+
+// Alignment of each entry in the archive entry table.
+// Padding between entries in the table is zero filled.
+#define IREE_IO_PARAMETER_ARCHIVE_HEADER_ALIGNMENT 16
+
+// Alignment of each entry in the archive entry table.
+// Padding between entries in the table is zero filled.
+#define IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT 16
+
+// Default alignment used for entry data storage.
+#define IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT 64
+
+// Parameter archive header magic identifier.
+// "IREE Parameter Archive"
+// "IRPA" = 0x49 0x52 0x50 0x41
+typedef uint32_t iree_io_parameter_archive_magic_t;
+#define IREE_IO_PARAMETER_ARCHIVE_MAGIC 0x41505249u
+
+// Bitmask of flags denoting archive traits.
+typedef uint64_t iree_io_parameter_archive_flags_t;
+
+// A range of bytes relative to an offset as defined in the header.
+typedef struct iree_io_parameter_archive_range_t {
+ iree_io_physical_offset_t offset;
+ iree_io_physical_size_t length;
+} iree_io_parameter_archive_range_t;
+// A range of bytes relative to the the header metadata_segment_offset.
+typedef iree_io_parameter_archive_range_t
+ iree_io_parameter_archive_metadata_ref_t;
+// A range of bytes relative to the the header storage_segment_offset.
+typedef iree_io_parameter_archive_range_t
+ iree_io_parameter_archive_storage_ref_t;
+
+// Header denoting the start of the of the archive in the enclosing file.
+// Needs not start at file offset zero but must be aligned to a multiple of
+// IREE_IO_PARAMETER_ARCHIVE_HEADER_ALIGNMENT. A single file may contain
+// multiple archives by using the `next_header_offset` field to build a linked
+// list of archive headers. All offsets in the header are relative to the header
+// offset so a global offset is not required even with multiple archives.
+//
+// NOTE: the following fields must always exist in all header versions as
+// parsers must be able to verify a file is an archive and check its version
+// to emit nice errors. Parsers are allowed to skip headers they can't load so
+// we also want to ensure the header size and any linked headers are available
+// regardless of version.
+typedef struct iree_io_parameter_archive_header_prefix_t {
+ // Magic header bytes; must be IREE_IO_PARAMETER_ARCHIVE_MAGIC.
+ iree_io_parameter_archive_magic_t magic;
+
+ // Major version identifier; major versions are incompatible.
+ uint16_t version_major;
+ // Minor version identifier; new minor versions with the same major are
+ // forward-compatible but may not be backwards compatible.
+ uint16_t version_minor;
+
+ // Total header size, including the 4 magic bytes.
+ iree_io_physical_size_t header_size;
+
+ // Optional (if 0) offset from the base of the header to another
+ // iree_io_parameter_archive_header_t in the file. This allows files to be
+ // concatenated or split without needing to re-encode any headers.
+ iree_io_physical_offset_t next_header_offset;
+
+ // Reserved for future use.
+ iree_io_parameter_archive_flags_t flags;
+} iree_io_parameter_archive_header_prefix_t;
+
+// Archive header for version_major 0.
+typedef struct iree_io_parameter_archive_header_v0_t {
+ // Common header prefix.
+ iree_io_parameter_archive_header_prefix_t prefix;
+
+ // Total number of entries in the entry table.
+ iree_io_physical_size_t entry_count;
+
+ // Offset from the base of the header to where the entry table resides.
+ // May overlap with other segments. Each entry in the table is aligned to
+ // IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT.
+ iree_io_parameter_archive_range_t entry_segment;
+
+ // Offset from the base of the header to where the metadata segment resides.
+ // May overlap with other segments.
+ iree_io_parameter_archive_range_t metadata_segment;
+
+ // Offset from the base of the header to where embedded data storage resides.
+ // May overlap with other segments.
+ iree_io_parameter_archive_range_t storage_segment;
+} iree_io_parameter_archive_header_v0_t;
+
+enum iree_io_parameter_archive_entry_type_e {
+ // Entry is ignored during processing as if it didn't exist in the table.
+ // Tools can erase entries by changing their type to this value.
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SKIP = 0,
+ // Entry represents a repeating pattern splatted into memory.
+ // See iree_io_parameter_archive_splat_entry_t.
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SPLAT = 1,
+ // Entry represents data embedded in the archive.
+ // See iree_io_parameter_archive_data_entry_t.
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_DATA = 2,
+ // Entry represents data stored in an external file.
+ // See iree_io_parameter_archive_external_entry_t.
+ IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_EXTERNAL = 3,
+};
+// Defines the type of an entry in the archive entry table.
+typedef uint32_t iree_io_parameter_archive_entry_type_t;
+
+// Entry type-specific flag bits.
+typedef uint64_t iree_io_parameter_archive_entry_flags_t;
+
+// Header shared across all entry types.
+// The total size of the header and the entry are included to allow for
+// traversing archives where not all entry types are known by the reader.
+// Each entry must be aligned to IREE_IO_PARAMETER_ARCHIVE_ENTRY_ALIGNMENT.
+// Padding between entries in the table are zero-filled.
+//
+// Entry names are arbitrary strings (could be paths, nested identifiers, etc).
+// Entry metadata is not defined by this format.
+typedef struct iree_io_parameter_archive_entry_header_t {
+ // Total size of the entry in the table, including this header but excluding
+ // any trailing padding.
+ iree_io_physical_size_t entry_size;
+ // Type of the entry indicating the outer structure containing this header.
+ iree_io_parameter_archive_entry_type_t type;
+ // For use by each entry type.
+ iree_io_parameter_archive_entry_flags_t flags;
+ // Reference to a non-NUL-terminated entry name in the metadata segment.
+ iree_io_parameter_archive_metadata_ref_t name;
+ // Reference to a metadata blob in the metadata segment.
+ // If the metadata is a string it need not have a NUL terminator.
+ // A zero length indicates no metadata is present.
+ iree_io_parameter_archive_metadata_ref_t metadata;
+ // Minimum alignment required when stored in the parent file.
+ // Implementations may require the absolute offset in the file to satisfy this
+ // alignment. Zero if alignment is unspecified.
+ iree_io_physical_size_t minimum_alignment;
+} iree_io_parameter_archive_entry_header_t;
+
+// An entry that contains a repeating pattern in virtual memory.
+// The pattern length must evenly divide the entry length and be a power of
+// two value of 1, 2, 4, 8, or 16 to allow for up to complex128 (complex<f64>).
+typedef struct iree_io_parameter_archive_splat_entry_t {
+ // Entry header with type IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_SPLAT.
+ iree_io_parameter_archive_entry_header_t header;
+ // Total length of the splat parameter in bytes.
+ iree_io_physical_size_t length;
+ // Little-endian pattern used to fill the range in memory.
+ // Examples:
+ // 0xAA: pattern_length=1 pattern=[0xAA ...]
+ // 0xBBAA: pattern_length=2 pattern=[0xAA 0xBB ...]
+ // 0xDDCCBBAA: pattern_length=4 pattern=[0xAA 0xBB 0xCC 0xDD ...]
+ uint8_t pattern[16];
+ // Length of the pattern in bytes defining how many bytes of the pattern
+ // field are valid from index 0.
+ uint8_t pattern_length;
+} iree_io_parameter_archive_splat_entry_t;
+
+// An entry referencing a span of data in the archive data storage segment.
+// Multiple entries with distinct virtual ranges may reference the same or
+// overlapping physical ranges.
+typedef struct iree_io_parameter_archive_data_entry_t {
+ // Entry header with type IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_DATA.
+ iree_io_parameter_archive_entry_header_t header;
+ // Relative offset of the data in the data storage segment.
+ // Total length of the data in the data storage segment. If less than the
+ // virtual length of the entry then it will be padded with zeros when loaded.
+ iree_io_parameter_archive_storage_ref_t storage;
+} iree_io_parameter_archive_data_entry_t;
+
+// An entry referencing data in an external file.
+typedef struct iree_io_parameter_archive_external_entry_t {
+ // Entry header with type IREE_IO_PARAMETER_ARCHIVE_ENTRY_TYPE_EXTERNAL.
+ iree_io_parameter_archive_entry_header_t header;
+ // Reference to a non-NUL-terminated file path in the metadata segment.
+ iree_io_parameter_archive_metadata_ref_t path;
+ // Absolute offset and length of the data in the external file.
+ iree_io_parameter_archive_range_t range;
+} iree_io_parameter_archive_external_entry_t;
+
+IREE_IO_PACKED_END
+
+#endif // IREE_SCHEMAS_PARAMETER_ARCHIVE_H_
diff --git a/runtime/src/iree/tooling/BUILD.bazel b/runtime/src/iree/tooling/BUILD.bazel
index 55d4d3c..5d3f676 100644
--- a/runtime/src/iree/tooling/BUILD.bazel
+++ b/runtime/src/iree/tooling/BUILD.bazel
@@ -161,6 +161,7 @@
"//runtime/src/iree/io:parameter_provider",
"//runtime/src/iree/io:scope_map",
"//runtime/src/iree/io/formats/gguf",
+ "//runtime/src/iree/io/formats/irpa",
"//runtime/src/iree/io/formats/safetensors",
"//runtime/src/iree/modules/io/parameters",
"//runtime/src/iree/vm",
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index a8d0711..0b5084a 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -180,6 +180,7 @@
iree::base::internal::path
iree::hal
iree::io::formats::gguf
+ iree::io::formats::irpa
iree::io::formats::safetensors
iree::io::parameter_index
iree::io::parameter_index_provider
diff --git a/runtime/src/iree/tooling/parameter_util.c b/runtime/src/iree/tooling/parameter_util.c
index 2708c08..fad5c72 100644
--- a/runtime/src/iree/tooling/parameter_util.c
+++ b/runtime/src/iree/tooling/parameter_util.c
@@ -9,8 +9,9 @@
#include "iree/base/internal/file_io.h"
#include "iree/base/internal/flags.h"
#include "iree/base/internal/path.h"
-#include "iree/io/formats/gguf/gguf_format.h"
-#include "iree/io/formats/safetensors/safetensors_format.h"
+#include "iree/io/formats/gguf/gguf_parser.h"
+#include "iree/io/formats/irpa/irpa_parser.h"
+#include "iree/io/formats/safetensors/safetensors_parser.h"
#include "iree/io/parameter_index.h"
#include "iree/io/parameter_index_provider.h"
#include "iree/io/scope_map.h"
@@ -110,7 +111,9 @@
// Index the file based on its (inferred) format.
iree_status_t status = iree_ok_status();
iree_string_view_t path_ext = iree_file_path_extension(path);
- if (iree_string_view_equal_case(path_ext, IREE_SV("gguf"))) {
+ if (iree_string_view_equal_case(path_ext, IREE_SV("irpa"))) {
+ status = iree_io_parse_irpa_index(file_handle, index);
+ } else if (iree_string_view_equal_case(path_ext, IREE_SV("gguf"))) {
status = iree_io_parse_gguf_index(file_handle, index);
} else if (iree_string_view_equal_case(path_ext, IREE_SV("safetensors"))) {
status = iree_io_parse_safetensors_index(file_handle, index);
diff --git a/runtime/src/iree/tooling/parameter_util.h b/runtime/src/iree/tooling/parameter_util.h
index 5342b89..a0e46c1 100644
--- a/runtime/src/iree/tooling/parameter_util.h
+++ b/runtime/src/iree/tooling/parameter_util.h
@@ -14,6 +14,7 @@
extern "C" {
#endif // __cplusplus
+typedef struct iree_io_parameter_index_t iree_io_parameter_index_t;
typedef struct iree_io_scope_map_t iree_io_scope_map_t;
// Populates |scope_map| with parameter indices as specified by flags.
diff --git a/tools/BUILD.bazel b/tools/BUILD.bazel
index d6316c1..833919c 100644
--- a/tools/BUILD.bazel
+++ b/tools/BUILD.bazel
@@ -80,6 +80,21 @@
)
iree_runtime_cc_binary(
+ name = "iree-convert-parameters",
+ srcs = ["iree-convert-parameters-main.c"],
+ deps = [
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/internal:file_io",
+ "//runtime/src/iree/base/internal:flags",
+ "//runtime/src/iree/hal",
+ "//runtime/src/iree/io:parameter_index",
+ "//runtime/src/iree/io:scope_map",
+ "//runtime/src/iree/io/formats/irpa",
+ "//runtime/src/iree/tooling:parameter_util",
+ ],
+)
+
+iree_runtime_cc_binary(
name = "iree-cpuinfo",
srcs = ["iree-cpuinfo.c"],
deps = [
@@ -90,6 +105,21 @@
)
iree_runtime_cc_binary(
+ name = "iree-create-parameters",
+ srcs = ["iree-create-parameters-main.c"],
+ deps = [
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/internal:file_io",
+ "//runtime/src/iree/base/internal:flags",
+ "//runtime/src/iree/hal",
+ "//runtime/src/iree/io:parameter_index",
+ "//runtime/src/iree/io:scope_map",
+ "//runtime/src/iree/io:stream",
+ "//runtime/src/iree/io/formats/irpa",
+ ],
+)
+
+iree_runtime_cc_binary(
name = "iree-dump-instruments",
srcs = ["iree-dump-instruments-main.c"],
deps = [
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index 65608e3..84c1a85 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -111,6 +111,22 @@
iree_cc_binary(
NAME
+ iree-convert-parameters
+ SRCS
+ "iree-convert-parameters-main.c"
+ DEPS
+ iree::base
+ iree::base::internal::file_io
+ iree::base::internal::flags
+ iree::hal
+ iree::io::formats::irpa
+ iree::io::parameter_index
+ iree::io::scope_map
+ iree::tooling::parameter_util
+)
+
+iree_cc_binary(
+ NAME
iree-cpuinfo
SRCS
"iree-cpuinfo.c"
@@ -122,6 +138,22 @@
iree_cc_binary(
NAME
+ iree-create-parameters
+ SRCS
+ "iree-create-parameters-main.c"
+ DEPS
+ iree::base
+ iree::base::internal::file_io
+ iree::base::internal::flags
+ iree::hal
+ iree::io::formats::irpa
+ iree::io::parameter_index
+ iree::io::scope_map
+ iree::io::stream
+)
+
+iree_cc_binary(
+ NAME
iree-dump-instruments
SRCS
"iree-dump-instruments-main.c"
diff --git a/tools/iree-convert-parameters-main.c b/tools/iree-convert-parameters-main.c
new file mode 100644
index 0000000..72b5a2d
--- /dev/null
+++ b/tools/iree-convert-parameters-main.c
@@ -0,0 +1,292 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+// Converts one or more parameter files into a single IREE Parameter Archive.
+// Allows for stripping and renaming parameters as basic editing features.
+
+#include <ctype.h>
+#include <stdio.h>
+
+#include "iree/base/api.h"
+#include "iree/base/internal/file_io.h"
+#include "iree/base/internal/flags.h"
+#include "iree/hal/api.h"
+#include "iree/io/formats/irpa/irpa_builder.h"
+#include "iree/io/parameter_index.h"
+#include "iree/io/scope_map.h"
+#include "iree/tooling/parameter_util.h"
+
+//===----------------------------------------------------------------------===//
+// Parameter index logic
+//===----------------------------------------------------------------------===//
+
+IREE_FLAG_LIST(string, exclude,
+ "Excludes a named parameter from the resulting file.");
+IREE_FLAG_LIST(string, rename,
+ "Renames a parameter when adding to the resulting file in the "
+ "form of `--rename=old=new`.");
+IREE_FLAG(bool, strip, false,
+ "Strips all parameters by replacing them with zeros.");
+IREE_FLAG_LIST(
+ string, splat,
+ "Turns a parameter into a splat of 0 (`--splat=name`) or a specific\n"
+ "sequence of typed values (`--splat=name=i8=123`, `--splat=name=f32=4.5`,\n"
+ "`--splat=name=x32=CAFEF00D`).");
+
+static bool iree_tooling_is_parameter_excluded(iree_string_view_t name) {
+ for (iree_host_size_t i = 0; i < FLAG_exclude_list().count; ++i) {
+ if (iree_string_view_equal(FLAG_exclude_list().values[i], name)) {
+ return true;
+ }
+ }
+ return false;
+}
+
+static iree_string_view_t iree_tooling_get_renamed_parameter_name(
+ iree_string_view_t name) {
+ for (iree_host_size_t i = 0; i < FLAG_rename_list().count; ++i) {
+ iree_string_view_t old_name, new_name;
+ iree_string_view_split(FLAG_rename_list().values[i], '=', &old_name,
+ &new_name);
+ if (iree_string_view_equal(old_name, name)) {
+ return new_name;
+ }
+ }
+ return name;
+}
+
+// Expects `type=value` consistent with the HAL.
+static iree_status_t iree_tooling_parse_splat(iree_string_view_t splat_value,
+ uint8_t* out_pattern_length,
+ uint8_t* out_pattern) {
+ if (iree_string_view_is_empty(splat_value)) {
+ *out_pattern_length = 1;
+ out_pattern[0] = 0;
+ return iree_ok_status();
+ }
+
+ iree_string_view_t type_str, value_str;
+ iree_string_view_split(splat_value, '=', &type_str, &value_str);
+
+ iree_hal_element_type_t type = IREE_HAL_ELEMENT_TYPE_NONE;
+ IREE_RETURN_IF_ERROR(iree_hal_parse_element_type(type_str, &type));
+
+ iree_device_size_t byte_count = iree_hal_element_dense_byte_count(type);
+ if (byte_count > 16) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "element type size for %.*s out of range of splat patterns",
+ (int)type_str.size, type_str.data);
+ }
+ *out_pattern_length = (uint8_t)byte_count;
+
+ return iree_hal_parse_element(value_str, type,
+ iree_make_byte_span(out_pattern, 16));
+}
+
+static iree_status_t iree_tooling_replace_splatted_parameter(
+ iree_io_parameter_index_entry_t* entry) {
+ // Always favor specific splat values.
+ for (iree_host_size_t i = 0; i < FLAG_splat_list().count; ++i) {
+ iree_string_view_t name, splat_value;
+ iree_string_view_split(FLAG_splat_list().values[i], '=', &name,
+ &splat_value);
+ if (iree_string_view_equal(name, entry->key)) {
+ entry->type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT;
+ memset(&entry->storage, 0, sizeof(entry->storage));
+ return iree_tooling_parse_splat(splat_value,
+ &entry->storage.splat.pattern_length,
+ entry->storage.splat.pattern);
+ }
+ }
+
+ // If not specifically splatted then see if we are stripping and use that.
+ if (FLAG_strip) {
+ entry->type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT;
+ memset(&entry->storage, 0, sizeof(entry->storage));
+ entry->storage.splat.pattern_length = 1;
+ entry->storage.splat.pattern[0] = 0;
+ return iree_ok_status();
+ }
+
+ return iree_ok_status();
+}
+
+static iree_status_t iree_tooling_convert_parameter_index(
+ iree_io_parameter_index_t* source_index,
+ iree_io_parameter_index_t* target_index) {
+ for (iree_host_size_t i = 0; i < iree_io_parameter_index_count(source_index);
+ ++i) {
+ // Get the existing entry we'll use as a template.
+ const iree_io_parameter_index_entry_t* source_entry = NULL;
+ IREE_RETURN_IF_ERROR(
+ iree_io_parameter_index_get(source_index, i, &source_entry));
+ iree_io_parameter_index_entry_t target_entry = *source_entry;
+
+ // If the parameter is in the exclude list then we just skip it.
+ if (iree_tooling_is_parameter_excluded(source_entry->key)) continue;
+
+ // If the parameter is in the rename list we'll add it with the new name.
+ target_entry.key =
+ iree_tooling_get_renamed_parameter_name(source_entry->key);
+
+ // If the parameter is turned into a splat we change its type. Note that it
+ // may have already been a splat but the user may want to change the value.
+ IREE_RETURN_IF_ERROR(
+ iree_tooling_replace_splatted_parameter(&target_entry));
+
+ // Add the entry (potentially modified) to the new index.
+ IREE_RETURN_IF_ERROR(
+ iree_io_parameter_index_add(target_index, &target_entry));
+ }
+ return iree_ok_status();
+}
+
+static iree_status_t iree_tooling_convert_parameters(
+ iree_io_scope_map_t* scope_map, iree_io_parameter_index_t* target_index,
+ iree_allocator_t host_allocator) {
+ for (iree_host_size_t i = 0; i < scope_map->count; ++i) {
+ IREE_RETURN_IF_ERROR(iree_tooling_convert_parameter_index(
+ scope_map->entries[i]->index, target_index));
+ }
+ return iree_ok_status();
+}
+
+//===----------------------------------------------------------------------===//
+// main
+//===----------------------------------------------------------------------===//
+
+IREE_FLAG(bool, quiet, false,
+ "Silences additional stdout output when not needed.");
+
+IREE_FLAG(string, output, "", "Output .irpa file path.");
+
+static void iree_io_file_handle_release_mapping(
+ void* user_data, iree_io_file_handle_primitive_t handle_primitive) {
+ iree_file_contents_free((iree_file_contents_t*)user_data);
+}
+
+typedef struct {
+ iree_allocator_t host_allocator;
+ const char* path;
+} iree_tooling_open_params_t;
+static iree_status_t iree_tooling_open_output_parameter_file(
+ void* user_data, iree_io_physical_offset_t archive_offset,
+ iree_io_physical_size_t archive_length,
+ iree_io_file_handle_t** out_file_handle) {
+ iree_tooling_open_params_t* params = (iree_tooling_open_params_t*)user_data;
+ iree_file_contents_t* file_contents = NULL;
+ IREE_RETURN_IF_ERROR(
+ iree_file_create_mapped(params->path, archive_offset + archive_length,
+ archive_offset, (iree_host_size_t)archive_length,
+ params->host_allocator, &file_contents));
+ iree_io_file_handle_release_callback_t release_callback = {
+ .fn = iree_io_file_handle_release_mapping,
+ .user_data = file_contents,
+ };
+ iree_status_t status = iree_io_file_handle_wrap_host_allocation(
+ IREE_IO_FILE_ACCESS_WRITE, file_contents->buffer, release_callback,
+ params->host_allocator, out_file_handle);
+ if (!iree_status_is_ok(status)) {
+ iree_file_contents_free(file_contents);
+ }
+ return status;
+}
+
+int main(int argc, char** argv) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_allocator_t host_allocator = iree_allocator_system();
+ int exit_code = EXIT_SUCCESS;
+
+ // Parse command line flags.
+ iree_flags_set_usage(
+ "iree-convert-parameters",
+ "Converts supported parameter file formats into IREE Parameter Archives\n"
+ "(.irpa) files. Provide one or more input parameter files in the same\n"
+ "form as expected by the iree-run-module tool (`--parameters=foo.gguf`)\n"
+ "and an output file with `--output=file.irpa`.\n"
+ "\n"
+ "Example converting from safetensors to IRPA:\n"
+ " iree-convert-parameters \\\n"
+ " --parameters=input.safetensors \\\n"
+ " --output=output.irpa\n"
+ "\n"
+ "Example mutating parameters:\n"
+ " iree-convert-parameters \\\n"
+ " --parameters=a.gguf \\\n"
+ " --parameters=b.safetensors \\\n"
+ " --exclude=unneeded_param \\\n"
+ " --rename=old_name=new_name \\\n"
+ " --splat=some_name=f32=4.2 \\\n"
+ " --output=ab.irpa\n"
+ "\n"
+ "Example stripping all parameters and replacing them with zeros except\n"
+ "for one that needs special handling:\n"
+ " iree-convert-parameters \\\n"
+ " --parameters=input.irpa \\\n"
+ " --strip \\\n"
+ " --splat=special_param=f32=1.0 \\\n"
+ " --output=output.irpa\n");
+ iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
+
+ // Load parameter indices as specified by command line flags.
+ iree_io_scope_map_t scope_map = {0};
+ iree_io_scope_map_initialize(host_allocator, &scope_map);
+ iree_status_t status =
+ iree_tooling_build_parameter_indices_from_flags(&scope_map);
+
+ // Build the new combined/modified index in memory based on the inputs.
+ iree_io_parameter_index_t* new_index = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_io_parameter_index_create(host_allocator, &new_index);
+ }
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_tooling_convert_parameters(&scope_map, new_index, host_allocator);
+ }
+ iree_io_scope_map_deinitialize(&scope_map);
+
+ iree_io_parameter_index_t* built_index = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_io_parameter_index_create(host_allocator, &built_index);
+ }
+
+ // Write out the new archive.
+ if (iree_status_is_ok(status)) {
+ iree_tooling_open_params_t open_params = {
+ .host_allocator = host_allocator,
+ .path = FLAG_output,
+ };
+ iree_io_parameter_archive_file_open_callback_t open_callback = {
+ .fn = iree_tooling_open_output_parameter_file,
+ .user_data = &open_params,
+ };
+ status = iree_io_build_parameter_archive(
+ new_index, built_index, open_callback,
+ /*target_file_offset=*/0, host_allocator);
+ }
+
+ // Dump the new index ala iree-dump-parameters to show the final file.
+ if (iree_status_is_ok(status) && !FLAG_quiet) {
+ status = iree_io_parameter_index_fprint(stdout, iree_string_view_empty(),
+ built_index);
+ }
+
+ iree_io_parameter_index_release(built_index);
+ iree_io_parameter_index_release(new_index);
+
+ fflush(stdout);
+ if (!iree_status_is_ok(status)) {
+ iree_status_fprint(stderr, status);
+ iree_status_free(status);
+ exit_code = EXIT_FAILURE;
+ }
+ fflush(stderr);
+
+ IREE_TRACE_ZONE_END(z0);
+ return exit_code;
+}
diff --git a/tools/iree-create-parameters-main.c b/tools/iree-create-parameters-main.c
new file mode 100644
index 0000000..81be97d
--- /dev/null
+++ b/tools/iree-create-parameters-main.c
@@ -0,0 +1,326 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+// Converts one or more parameter files into a single IREE Parameter Archive.
+// Allows for stripping and renaming parameters as basic editing features.
+
+#include <ctype.h>
+#include <stdio.h>
+
+#include "iree/base/api.h"
+#include "iree/base/internal/file_io.h"
+#include "iree/base/internal/flags.h"
+#include "iree/hal/api.h"
+#include "iree/io/formats/irpa/irpa_builder.h"
+#include "iree/io/parameter_index.h"
+#include "iree/io/scope_map.h"
+#include "iree/io/stream.h"
+
+//===----------------------------------------------------------------------===//
+// Parameter builder logic
+//===----------------------------------------------------------------------===//
+
+IREE_FLAG_LIST(
+ string, splat,
+ "Declares a splat parameter in the form `name=shapextype=pattern`.\n"
+ "Splat parameters have no storage on disk and can be used to mock real\n"
+ "parameters that besides load-time performance and contents behave the\n"
+ "same as if real parameters had been used. The shape and type are used to\n"
+ "calculate the parameter size in the index and the value is interpreted\n"
+ "based on the type. Note that splat parameters cannot be mutated at\n"
+ "runtime and can only be used for constant values.\n");
+
+// TODO(benvanik): support @file.bin syntax, numpy files, etc.
+IREE_FLAG_LIST(
+ string, data,
+ "Declares a data parameter in the form `name=shapextype=pattern`.\n"
+ "Data parameters have a storage reservation in the final archive with\n"
+ "either zeroed contents or a specified repeating pattern value.\n"
+ "The shape and type are used to calculate the parameter size in the index\n"
+ "and the value is interpreted based on the type. Omitting the value\n"
+ "will leave the parameter with zeroed contents on disk.");
+
+// TODO(benvanik): support external entries (--external=file.bin#128-512).
+
+IREE_FLAG(int32_t, alignment, IREE_IO_PARAMETER_ARCHIVE_DEFAULT_DATA_ALIGNMENT,
+ "Storage data alignment relative to the header.");
+
+typedef struct {
+ iree_string_view_t name;
+ uint64_t storage_size;
+ uint64_t element_count;
+ iree_hal_element_type_t element_type;
+ struct {
+ uint8_t pattern[16];
+ uint8_t pattern_length;
+ } splat;
+} iree_io_parameter_info_t;
+
+// Parses a `name=shapextype[=value]` flag string.
+static iree_status_t iree_io_parameter_info_from_string(
+ iree_string_view_t parameter_value, iree_io_parameter_info_t* out_info) {
+ memset(out_info, 0, sizeof(*out_info));
+
+ iree_string_view_t spec;
+ iree_string_view_split(parameter_value, '=', &out_info->name, &spec);
+
+ iree_string_view_t shape_type, contents;
+ iree_string_view_split(spec, '=', &shape_type, &contents);
+
+ iree_host_size_t shape_rank = 0;
+ iree_hal_dim_t shape[16] = {0};
+ IREE_RETURN_IF_ERROR(iree_hal_parse_shape_and_element_type(
+ shape_type, IREE_ARRAYSIZE(shape), &shape_rank, shape,
+ &out_info->element_type));
+
+ if (IREE_UNLIKELY(iree_hal_element_bit_count(out_info->element_type) == 0) ||
+ IREE_UNLIKELY(
+ !iree_hal_element_is_byte_aligned(out_info->element_type))) {
+ return iree_make_status(
+ IREE_STATUS_INVALID_ARGUMENT,
+ "opaque and sub-byte aligned element types cannot currently be used as "
+ "splats; use a bitcast type (2xi4=1xi8, etc)");
+ }
+
+ out_info->storage_size =
+ iree_hal_element_dense_byte_count(out_info->element_type);
+ out_info->element_count = 1;
+ for (iree_host_size_t i = 0; i < shape_rank; ++i) {
+ out_info->storage_size *= shape[i];
+ out_info->element_count *= shape[i];
+ }
+
+ // TODO(benvanik): support external files and such; for now we just assume
+ // either empty string (0 splat) or a splat value.
+
+ if (iree_string_view_is_empty(contents)) {
+ out_info->splat.pattern_length = 1;
+ memset(out_info->splat.pattern, 0, sizeof(out_info->splat.pattern));
+ } else {
+ iree_device_size_t byte_count =
+ iree_hal_element_dense_byte_count(out_info->element_type);
+ if (byte_count > sizeof(out_info->splat.pattern)) {
+ return iree_make_status(
+ IREE_STATUS_OUT_OF_RANGE,
+ "element type size for %.*s out of range of splat patterns",
+ (int)shape_type.size, shape_type.data);
+ }
+ out_info->splat.pattern_length = (uint8_t)byte_count;
+ IREE_RETURN_IF_ERROR(iree_hal_parse_element(
+ contents, out_info->element_type,
+ iree_make_byte_span(out_info->splat.pattern, 16)));
+ }
+ return iree_ok_status();
+}
+
+// Declares parameter metadata for all parameters specified by flags.
+static iree_status_t iree_tooling_declare_parameters(
+ iree_io_parameter_archive_builder_t* builder) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ // Metadata-only parameters first; they have no storage and the fact that they
+ // are not in-order with any data parameters means they don't impact runtime
+ // access locality behavior.
+ for (iree_host_size_t i = 0; i < FLAG_splat_list().count; ++i) {
+ iree_io_parameter_info_t info;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_parameter_info_from_string(FLAG_splat_list().values[i], &info));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_parameter_archive_builder_add_splat_entry(
+ builder, info.name, /*metadata=*/iree_const_byte_span_empty(),
+ info.splat.pattern, info.splat.pattern_length, info.storage_size));
+ }
+
+ // Data parameters follow and will appear in storage in the order they were
+ // declared with flags.
+ for (iree_host_size_t i = 0; i < FLAG_data_list().count; ++i) {
+ iree_io_parameter_info_t info;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_parameter_info_from_string(FLAG_data_list().values[i], &info));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_parameter_archive_builder_add_data_entry(
+ builder, info.name, /*metadata=*/iree_const_byte_span_empty(),
+ FLAG_alignment, info.storage_size));
+ }
+
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+// Defines parameter storage for those that require it.
+static iree_status_t iree_tooling_define_parameters(
+ iree_io_parameter_index_t* target_index,
+ iree_io_physical_offset_t target_file_offset,
+ iree_io_stream_t* target_stream) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ for (iree_host_size_t i = 0; i < FLAG_data_list().count; ++i) {
+ iree_io_parameter_info_t info;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_parameter_info_from_string(FLAG_data_list().values[i], &info));
+ const iree_io_parameter_index_entry_t* target_entry = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0,
+ iree_io_parameter_index_lookup(target_index, info.name, &target_entry));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_seek(
+ target_stream, IREE_IO_STREAM_SEEK_SET,
+ target_file_offset + target_entry->storage.file.offset));
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_io_stream_fill(target_stream, info.element_count,
+ info.splat.pattern, info.splat.pattern_length));
+ }
+
+ IREE_TRACE_ZONE_END(z0);
+ return iree_ok_status();
+}
+
+//===----------------------------------------------------------------------===//
+// main
+//===----------------------------------------------------------------------===//
+
+IREE_FLAG(bool, quiet, false,
+ "Silences additional stdout output when not needed.");
+
+// TODO(benvanik): add --append= to chain archive headers.
+IREE_FLAG(string, output, "", "Output .irpa file path.");
+
+static void iree_io_file_handle_release_mapping(
+ void* user_data, iree_io_file_handle_primitive_t handle_primitive) {
+ iree_file_contents_free((iree_file_contents_t*)user_data);
+}
+
+static iree_status_t iree_tooling_open_output_parameter_file(
+ iree_io_physical_offset_t archive_offset,
+ iree_io_physical_size_t archive_length, iree_allocator_t host_allocator,
+ iree_io_file_handle_t** out_file_handle) {
+ iree_file_contents_t* file_contents = NULL;
+ IREE_RETURN_IF_ERROR(iree_file_create_mapped(
+ FLAG_output, archive_offset + archive_length, archive_offset,
+ (iree_host_size_t)archive_length, host_allocator, &file_contents));
+ iree_io_file_handle_release_callback_t release_callback = {
+ .fn = iree_io_file_handle_release_mapping,
+ .user_data = file_contents,
+ };
+ iree_status_t status = iree_io_file_handle_wrap_host_allocation(
+ IREE_IO_FILE_ACCESS_WRITE, file_contents->buffer, release_callback,
+ host_allocator, out_file_handle);
+ if (!iree_status_is_ok(status)) {
+ iree_file_contents_free(file_contents);
+ }
+ return status;
+}
+
+int main(int argc, char** argv) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ iree_allocator_t host_allocator = iree_allocator_system();
+ int exit_code = EXIT_SUCCESS;
+
+ // Parse command line flags.
+ iree_flags_set_usage(
+ "iree-create-parameters",
+ "Creates IREE Parameter Archive (.irpa) files. Provide zero or more\n"
+ "parameter value declarations and an output file with\n"
+ "`--output=file.irpa` to produce a new file with zeroed or patterned\n"
+ "contents.\n"
+ "\n"
+ "Parameter declarations take a shape and type in order to calculate the\n"
+ "required storage size of the parameter at rest and at runtime. The\n"
+ "shape and type need not match what the consuming program expects so\n"
+ "long as the storage size is equivalent; for example, if the program\n"
+ "expects a parameter of type `tensor<8x2xi4>` the parameter declaration\n"
+ "can be `8xi8`, `1xi64`, `2xf32`, etc.\n"
+ "\n"
+ "Example creating a file with two embedded data parameters that have\n"
+ "zeroed contents and one with a repeating pattern:\n"
+ " iree-create-parameters \\\n"
+ " --data=my.zeroed_param_1=4096xf32 \\\n"
+ " --data=my.zeroed_param_2=2x4096xi16 \\\n"
+ " --data=my.pattern_param_2=8x2xf32=2.1 \\\n"
+ " --output=output_with_storage.irpa\n"
+ "\n"
+ "Example creating a file with splatted values (no storage on disk):\n"
+ " iree-create-parameters \\\n"
+ " --splat=my.splat_param_1=4096xf32=4.1 \\\n"
+ " --splat=my.splat_param_2=2x4096xi16=123 \\\n"
+ " --output=output_without_storage.irpa\n");
+ iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
+
+ iree_io_parameter_archive_builder_t builder;
+ iree_io_parameter_archive_builder_initialize(host_allocator, &builder);
+
+ // Declare parameters based on flags, populating the builder with the metadata
+ // for each parameter without yet writing any data.
+ iree_status_t status = iree_tooling_declare_parameters(&builder);
+
+ // Open a file of sufficient size (now that we know it) for writing.
+ iree_io_physical_offset_t target_file_offset = 0;
+ iree_io_physical_offset_t archive_offset = iree_align_uint64(
+ target_file_offset, IREE_IO_PARAMETER_ARCHIVE_HEADER_ALIGNMENT);
+ iree_io_physical_size_t archive_length =
+ iree_io_parameter_archive_builder_total_size(&builder);
+ iree_io_file_handle_t* target_file_handle = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_tooling_open_output_parameter_file(
+ archive_offset, archive_length, host_allocator, &target_file_handle);
+ }
+
+ // Wrap the target file in a stream.
+ iree_io_stream_t* target_stream = NULL;
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_io_stream_open(IREE_IO_STREAM_MODE_WRITABLE, target_file_handle,
+ target_file_offset, host_allocator, &target_stream);
+ }
+
+ // Allocate an index we'll populate during building to allow us to get the
+ // storage ranges of non-metadata parameters.
+ iree_io_parameter_index_t* built_index = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_io_parameter_index_create(host_allocator, &built_index);
+ }
+
+ // Commit the archive header to the file and produce an index referencing it.
+ // This will allow us to know where to copy file contents.
+ if (iree_status_is_ok(status)) {
+ status = iree_io_parameter_archive_builder_write(
+ &builder, target_file_handle, target_file_offset, target_stream,
+ built_index);
+ }
+
+ // Define non-metadata-only parameters that use the data storage segment.
+ if (iree_status_is_ok(status)) {
+ status = iree_tooling_define_parameters(built_index, target_file_offset,
+ target_stream);
+ }
+
+ // Dump the new index ala iree-dump-parameters to show the final file.
+ if (iree_status_is_ok(status) && !FLAG_quiet) {
+ status = iree_io_parameter_index_fprint(stdout, iree_string_view_empty(),
+ built_index);
+ }
+
+ iree_io_stream_release(target_stream);
+ iree_io_file_handle_release(target_file_handle);
+ iree_io_parameter_archive_builder_deinitialize(&builder);
+ iree_io_parameter_index_release(built_index);
+
+ fflush(stdout);
+ if (!iree_status_is_ok(status)) {
+ iree_status_fprint(stderr, status);
+ iree_status_free(status);
+ exit_code = EXIT_FAILURE;
+ }
+ fflush(stderr);
+
+ IREE_TRACE_ZONE_END(z0);
+ return exit_code;
+}
diff --git a/tools/iree-dump-parameters-main.c b/tools/iree-dump-parameters-main.c
index 2e2898a..4c6f3d2 100644
--- a/tools/iree-dump-parameters-main.c
+++ b/tools/iree-dump-parameters-main.c
@@ -30,73 +30,6 @@
#include "iree/tooling/parameter_util.h"
//===----------------------------------------------------------------------===//
-// Parameter index information
-//===----------------------------------------------------------------------===//
-
-static iree_status_t iree_tooling_dump_parameter_index(
- iree_string_view_t scope, iree_io_parameter_index_t* index) {
- iree_host_size_t entry_count = iree_io_parameter_index_count(index);
- uint64_t total_bytes = 0;
- for (iree_host_size_t i = 0; i < entry_count; ++i) {
- const iree_io_parameter_index_entry_t* entry = NULL;
- IREE_RETURN_IF_ERROR(iree_io_parameter_index_get(index, i, &entry));
- total_bytes += entry->length;
- }
- fprintf(stdout,
- "//"
- "===-----------------------------------------------------------------"
- "---------------------------------------------===//\n");
- fprintf(stdout,
- "// Parameter scope `%.*s` (%" PRIhsz " entries, %" PRIu64
- " total bytes)\n",
- (int)scope.size, scope.data, entry_count, total_bytes);
- fprintf(stdout,
- "//===------------+------------------+------------------+------------"
- "-----------------------------------------------===//\n");
- fprintf(stdout,
- "// Start | End | Length | Key\n");
- fprintf(stdout,
- "//---------------+------------------+------------------+------------"
- "--------------------------------------------------//\n");
- for (iree_host_size_t i = 0; i < entry_count; ++i) {
- const iree_io_parameter_index_entry_t* entry = NULL;
- IREE_RETURN_IF_ERROR(iree_io_parameter_index_get(index, i, &entry));
- switch (entry->type) {
- case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT:
- fprintf(stdout,
- " - | - | %16" PRIu64
- " | `%.*s`\n",
- entry->length, (int)entry->key.size, entry->key.data);
- break;
- case IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE:
- fprintf(stdout,
- "%16" PRIu64 " | %16" PRIu64 " | %16" PRIu64 " | `%.*s`\n",
- entry->storage.file.offset,
- entry->storage.file.offset + entry->length, entry->length,
- (int)entry->key.size, entry->key.data);
- break;
- default:
- fprintf(stdout,
- " ? | ? | %16" PRIu64 " | `%.*s`\n",
- entry->length, (int)entry->key.size, entry->key.data);
- break;
- }
- }
- fprintf(stdout, "\n");
- return iree_ok_status();
-}
-
-static iree_status_t iree_tooling_dump_scope_map(
- iree_io_scope_map_t* scope_map) {
- for (iree_host_size_t i = 0; i < scope_map->count; ++i) {
- iree_string_view_t scope = scope_map->entries[i]->scope;
- iree_io_parameter_index_t* index = scope_map->entries[i]->index;
- IREE_RETURN_IF_ERROR(iree_tooling_dump_parameter_index(scope, index));
- }
- return iree_ok_status();
-}
-
-//===----------------------------------------------------------------------===//
// Parameter extraction
//===----------------------------------------------------------------------===//
@@ -196,7 +129,14 @@
// Dump parameter information.
if (iree_status_is_ok(status)) {
- status = iree_tooling_dump_scope_map(&scope_map);
+ iree_string_builder_t sb;
+ iree_string_builder_initialize(host_allocator, &sb);
+ status = iree_io_scope_map_dump(&scope_map, &sb);
+ if (iree_status_is_ok(status)) {
+ fprintf(stdout, "%.*s", (int)iree_string_builder_size(&sb),
+ iree_string_builder_buffer(&sb));
+ }
+ iree_string_builder_deinitialize(&sb);
}
// Extract parameters as requested, if any.
