Supporting file descriptors in iree_io_stream_open. (#19665)
diff --git a/runtime/src/iree/io/BUILD.bazel b/runtime/src/iree/io/BUILD.bazel
index 7d4c0e1..892989f 100644
--- a/runtime/src/iree/io/BUILD.bazel
+++ b/runtime/src/iree/io/BUILD.bazel
@@ -17,19 +17,6 @@
srcs = ["file_handle.c"],
hdrs = ["file_handle.h"],
deps = [
- ":memory_stream",
- ":stream",
- "//runtime/src/iree/base",
- "//runtime/src/iree/base/internal",
- ],
-)
-
-iree_runtime_cc_library(
- name = "memory_stream",
- srcs = ["memory_stream.c"],
- hdrs = ["memory_stream.h"],
- deps = [
- ":stream",
"//runtime/src/iree/base",
"//runtime/src/iree/base/internal",
],
@@ -39,7 +26,7 @@
name = "memory_stream_test",
srcs = ["memory_stream_test.cc"],
deps = [
- ":memory_stream",
+ ":stream",
"//runtime/src/iree/base",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
@@ -93,32 +80,21 @@
)
iree_runtime_cc_library(
- name = "stdio_stream",
- srcs = ["stdio_stream.c"],
- hdrs = ["stdio_stream.h"],
- deps = [
- ":stream",
- "//runtime/src/iree/base",
- "//runtime/src/iree/base/internal",
- ],
-)
-
-iree_runtime_cc_library(
name = "stream",
- srcs = ["stream.c"],
- hdrs = ["stream.h"],
- deps = [
- "//runtime/src/iree/base",
- "//runtime/src/iree/base/internal",
+ srcs = [
+ "memory_stream.c",
+ "stdio_stream.c",
+ "stream.c",
+ "vec_stream.c",
],
-)
-
-iree_runtime_cc_library(
- name = "vec_stream",
- srcs = ["vec_stream.c"],
- hdrs = ["vec_stream.h"],
+ hdrs = [
+ "memory_stream.h",
+ "stdio_stream.h",
+ "stream.h",
+ "vec_stream.h",
+ ],
deps = [
- ":stream",
+ ":file_handle",
"//runtime/src/iree/base",
"//runtime/src/iree/base/internal",
],
@@ -128,7 +104,7 @@
name = "vec_stream_test",
srcs = ["vec_stream_test.cc"],
deps = [
- ":vec_stream",
+ ":stream",
"//runtime/src/iree/base",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
diff --git a/runtime/src/iree/io/CMakeLists.txt b/runtime/src/iree/io/CMakeLists.txt
index cf4fd9f..b6ad802 100644
--- a/runtime/src/iree/io/CMakeLists.txt
+++ b/runtime/src/iree/io/CMakeLists.txt
@@ -18,22 +18,6 @@
SRCS
"file_handle.c"
DEPS
- ::memory_stream
- ::stream
- iree::base
- iree::base::internal
- PUBLIC
-)
-
-iree_cc_library(
- NAME
- memory_stream
- HDRS
- "memory_stream.h"
- SRCS
- "memory_stream.c"
- DEPS
- ::stream
iree::base
iree::base::internal
PUBLIC
@@ -45,7 +29,7 @@
SRCS
"memory_stream_test.cc"
DEPS
- ::memory_stream
+ ::stream
iree::base
iree::testing::gtest
iree::testing::gtest_main
@@ -111,40 +95,19 @@
iree_cc_library(
NAME
- stdio_stream
- HDRS
- "stdio_stream.h"
- SRCS
- "stdio_stream.c"
- DEPS
- ::stream
- iree::base
- iree::base::internal
- PUBLIC
-)
-
-iree_cc_library(
- NAME
stream
HDRS
+ "memory_stream.h"
+ "stdio_stream.h"
"stream.h"
- SRCS
- "stream.c"
- DEPS
- iree::base
- iree::base::internal
- PUBLIC
-)
-
-iree_cc_library(
- NAME
- vec_stream
- HDRS
"vec_stream.h"
SRCS
+ "memory_stream.c"
+ "stdio_stream.c"
+ "stream.c"
"vec_stream.c"
DEPS
- ::stream
+ ::file_handle
iree::base
iree::base::internal
PUBLIC
@@ -156,7 +119,7 @@
SRCS
"vec_stream_test.cc"
DEPS
- ::vec_stream
+ ::stream
iree::base
iree::testing::gtest
iree::testing::gtest_main
diff --git a/runtime/src/iree/io/file_handle.c b/runtime/src/iree/io/file_handle.c
index 9a6727b..22d9d80 100644
--- a/runtime/src/iree/io/file_handle.c
+++ b/runtime/src/iree/io/file_handle.c
@@ -7,7 +7,6 @@
#include "iree/io/file_handle.h"
#include "iree/base/internal/atomics.h"
-#include "iree/io/memory_stream.h"
#if IREE_FILE_IO_ENABLE
#if defined(IREE_PLATFORM_WINDOWS)
@@ -644,7 +643,7 @@
return iree_make_status(iree_status_code_from_win32_error(GetLastError()),
"failed to query file handle information");
}
- const uint64_t file_size = file_info.AllocationSize.QuadPart;
+ const uint64_t file_size = file_info.EndOfFile.QuadPart;
// Validate and adjust view size if needed.
iree_host_size_t adjusted_length = 0;
@@ -887,91 +886,3 @@
IREE_ASSERT_ARGUMENT(mapping);
return mapping->contents;
}
-
-//===----------------------------------------------------------------------===//
-// 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_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 908f1e7..ba2f81b 100644
--- a/runtime/src/iree/io/file_handle.h
+++ b/runtime/src/iree/io/file_handle.h
@@ -10,7 +10,6 @@
#include <stdint.h>
#include "iree/base/api.h"
-#include "iree/io/stream.h"
#ifdef __cplusplus
extern "C" {
@@ -308,28 +307,6 @@
IREE_API_EXPORT iree_byte_span_t
iree_io_file_mapping_contents_rw(iree_io_file_mapping_t* mapping);
-//===----------------------------------------------------------------------===//
-// iree_io_stream_t utilities
-//===----------------------------------------------------------------------===//
-
-// TODO(benvanik): remove/rework iree_io_stream_open so that it doesn't pull in
-// any implementations by putting callbacks on the file handle constructors.
-
-// 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/stdio_stream.c b/runtime/src/iree/io/stdio_stream.c
index b190585..b26362c 100644
--- a/runtime/src/iree/io/stdio_stream.c
+++ b/runtime/src/iree/io/stdio_stream.c
@@ -17,13 +17,21 @@
#define IREE_SET_BINARY_MODE(handle) _setmode(_fileno(handle), O_BINARY)
+#define iree_dup _dup
+#define iree_close _close
+
#define iree_fseek _fseeki64
#define iree_ftell _ftelli64
#else
+#include <unistd.h>
+
#define IREE_SET_BINARY_MODE(handle) ((void)0)
+#define iree_dup dup
+#define iree_close close
+
#if _FILE_OFFSET_BITS == 64 || _POSIX_C_SOURCE >= 200112L
#define iree_fseek fseeko
#define iree_ftell ftello
@@ -91,6 +99,52 @@
}
#if IREE_FILE_IO_ENABLE
+
+// Populates the |out_fopen_mode| string for use with fopen-like calls based on
+// the iree_io_stdio_stream_mode_t bitmap.
+//
+// NOTE: not all implementations support all mode flags and this may have
+// different behavior. We should paper over it here but don't today given the
+// limited usage of this and our intent to rewrite it all using
+// platform-optimal APIs instead of stdio.
+static void iree_io_map_stdio_fopen_mode(iree_io_stdio_stream_mode_t stdio_mode,
+ char out_fopen_mode[16]) {
+ memset(out_fopen_mode, 0, 16);
+
+ if (iree_all_bits_set(stdio_mode, IREE_IO_STDIO_STREAM_MODE_READ |
+ IREE_IO_STDIO_STREAM_MODE_WRITE |
+ IREE_IO_STDIO_STREAM_MODE_APPEND)) {
+ strcat(out_fopen_mode, "a+");
+ } else if (iree_all_bits_set(stdio_mode,
+ IREE_IO_STDIO_STREAM_MODE_READ |
+ IREE_IO_STDIO_STREAM_MODE_WRITE |
+ IREE_IO_STDIO_STREAM_MODE_DISCARD)) {
+ strcat(out_fopen_mode, "w+");
+ } else if (iree_all_bits_set(stdio_mode,
+ IREE_IO_STDIO_STREAM_MODE_READ |
+ IREE_IO_STDIO_STREAM_MODE_WRITE)) {
+ strcat(out_fopen_mode, "r+");
+ } else if (iree_all_bits_set(stdio_mode,
+ IREE_IO_STDIO_STREAM_MODE_WRITE |
+ IREE_IO_STDIO_STREAM_MODE_APPEND)) {
+ strcat(out_fopen_mode, "a");
+ } else if (iree_all_bits_set(stdio_mode, IREE_IO_STDIO_STREAM_MODE_WRITE)) {
+ strcat(out_fopen_mode, "w");
+ } else if (iree_all_bits_set(stdio_mode, IREE_IO_STDIO_STREAM_MODE_READ)) {
+ strcat(out_fopen_mode, "r");
+ }
+ if (iree_all_bits_set(stdio_mode, IREE_IO_STDIO_STREAM_MODE_WRITE) &&
+ !iree_all_bits_set(stdio_mode, IREE_IO_STDIO_STREAM_MODE_DISCARD)) {
+ // If writable and not discard then the file must not exist.
+ // TODO(benvanik): actually observe this; the C11 spec says `x` is supported
+ // but at least on MSVC's CRT it isn't. We can emulate this with stat and
+ // such but today we don't have any uses that require it.
+ // strcat(out_fopen_mode, "x");
+ }
+ // Force binary mode (avoid Windows CRLF expansion).
+ strcat(out_fopen_mode, "b");
+}
+
IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open(
iree_io_stdio_stream_mode_t mode, iree_string_view_t path,
iree_allocator_t host_allocator, iree_io_stream_t** out_stream) {
@@ -107,40 +161,8 @@
stream_mode |= IREE_IO_STREAM_MODE_WRITABLE;
}
- // NOTE: not all implementations support all mode flags and this may have
- // different behavior. We should paper over it here but don't today given the
- // limited usage of this and our intent to rewrite it all using
- // platform-optimal APIs instead of stdio.
char fopen_mode[16] = {0};
- if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_READ |
- IREE_IO_STDIO_STREAM_MODE_WRITE |
- IREE_IO_STDIO_STREAM_MODE_APPEND)) {
- strcat(fopen_mode, "a+");
- } else if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_READ |
- IREE_IO_STDIO_STREAM_MODE_WRITE |
- IREE_IO_STDIO_STREAM_MODE_DISCARD)) {
- strcat(fopen_mode, "w+");
- } else if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_READ |
- IREE_IO_STDIO_STREAM_MODE_WRITE)) {
- strcat(fopen_mode, "r+");
- } else if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_WRITE |
- IREE_IO_STDIO_STREAM_MODE_APPEND)) {
- strcat(fopen_mode, "a");
- } else if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_WRITE)) {
- strcat(fopen_mode, "w");
- } else if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_READ)) {
- strcat(fopen_mode, "r");
- }
- if (iree_all_bits_set(stream_mode, IREE_IO_STREAM_MODE_WRITABLE) &&
- !iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_DISCARD)) {
- // If writable and not discard then the file must not exist.
- // TODO(benvanik): actually observe this; the C11 spec says `x` is supported
- // but at least on MSVC's CRT it isn't. We can emulate this with stat and
- // such but today we don't have any uses that require it.
- // strcat(fopen_mode, "x");
- }
- // Force binary mode (avoid Windows CRLF expansion).
- strcat(fopen_mode, "b");
+ iree_io_map_stdio_fopen_mode(mode, fopen_mode);
// Since we stack alloc the path we want to keep it reasonable.
// We could heap allocate instead but a few thousand chars is quite long and
@@ -188,7 +210,70 @@
IREE_TRACE_ZONE_END(z0);
return status;
}
+
+IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open_fd(
+ iree_io_stdio_stream_mode_t mode, int fd, iree_allocator_t host_allocator,
+ iree_io_stream_t** out_stream) {
+ IREE_ASSERT_ARGUMENT(out_stream);
+ *out_stream = NULL;
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ // Duplicate the file descriptor so that we have our own copy of the seek
+ // position. The initial position will be preserved.
+ int dup_fd = iree_dup(fd);
+ if (dup_fd == -1) {
+ IREE_TRACE_ZONE_END(z0);
+ return iree_make_stdio_status(
+ "unable to duplicate file descriptor; possibly out of file descriptors "
+ "(see ulimit)");
+ }
+
+ iree_io_stream_mode_t stream_mode = IREE_IO_STREAM_MODE_SEEKABLE;
+ if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_READ)) {
+ stream_mode |= IREE_IO_STREAM_MODE_READABLE;
+ }
+ if (iree_all_bits_set(mode, IREE_IO_STDIO_STREAM_MODE_WRITE)) {
+ stream_mode |= IREE_IO_STREAM_MODE_WRITABLE;
+ }
+
+ char fopen_mode[16] = {0};
+ iree_io_map_stdio_fopen_mode(mode, fopen_mode);
+
+ // NOTE: after this point the file handle is associated with dup_fd and
+ // anything we do to it (like closing) will apply to the dup_fd.
+ iree_status_t status = iree_ok_status();
+ FILE* handle = fdopen(dup_fd, fopen_mode);
+ if (handle == NULL) {
+ status = iree_make_stdio_statusf(
+ "unable to open file descriptor with mode %d", mode);
+ }
+
+ // Ownership of the handle (and the dup_fd backing it) is transferred.
+ iree_io_stream_t* stream = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_io_stdio_stream_wrap(
+ stream_mode, handle, /*owns_handle=*/true, host_allocator, &stream);
+ }
+
+ if (iree_status_is_ok(status)) {
+ *out_stream = stream;
+ } else {
+ if (stream) {
+ // NOTE: closes the file handle/dup_fd.
+ iree_io_stream_release(stream);
+ } else if (handle) {
+ // NOTE: closes the dup_fd.
+ fclose(handle);
+ } else if (dup_fd > 0) {
+ iree_close(dup_fd);
+ }
+ }
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
#else
+
IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open(
iree_io_stdio_stream_mode_t mode, iree_string_view_t path,
iree_allocator_t host_allocator, iree_io_stream_t** out_stream) {
@@ -196,6 +281,15 @@
"file support has been compiled out of this binary; "
"set IREE_FILE_IO_ENABLE=1 to include it");
}
+
+IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open_fd(
+ iree_io_stdio_stream_mode_t mode, int fd, iree_allocator_t host_allocator,
+ iree_io_stream_t** out_stream) {
+ return iree_make_status(IREE_STATUS_UNAVAILABLE,
+ "file support has been compiled out of this binary; "
+ "set IREE_FILE_IO_ENABLE=1 to include it");
+}
+
#endif // IREE_FILE_IO_ENABLE
static void iree_io_stdio_stream_destroy(
@@ -297,16 +391,16 @@
// reads in one go anyway.
iree_host_size_t bytes_read = 0;
while (bytes_read < buffer_capacity) {
- iree_host_size_t chunk_size =
+ const iree_host_size_t chunk_size =
iree_min(buffer_capacity - bytes_read, INT_MAX);
- iree_host_size_t read_size =
+ const iree_host_size_t read_size =
fread((uint8_t*)buffer + bytes_read, 1, chunk_size, stream->handle);
if (read_size != chunk_size) {
// Failed to read chunk - may have reached EOF.
if (feof(stream->handle)) {
if (out_buffer_length) {
// Ok to hit EOF; just return what's valid.
- *out_buffer_length = bytes_read + read_size;
+ bytes_read += read_size;
} else {
status = iree_make_status(IREE_STATUS_OUT_OF_RANGE,
"end-of-file encountered during read");
@@ -319,6 +413,9 @@
bytes_read += read_size;
}
+ if (out_buffer_length) {
+ *out_buffer_length = bytes_read;
+ }
IREE_TRACE_ZONE_END(z0);
return status;
}
diff --git a/runtime/src/iree/io/stdio_stream.h b/runtime/src/iree/io/stdio_stream.h
index ddb2773..aef08ef 100644
--- a/runtime/src/iree/io/stdio_stream.h
+++ b/runtime/src/iree/io/stdio_stream.h
@@ -51,6 +51,13 @@
iree_io_stdio_stream_mode_t mode, iree_string_view_t path,
iree_allocator_t host_allocator, iree_io_stream_t** out_stream);
+// Opens a file descriptor |fd| using fdopen with the mode determined by |mode|.
+// The file descriptor is duplicated and seeking the stream will not be visible
+// with any other users of the file descriptor.
+IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open_fd(
+ iree_io_stdio_stream_mode_t mode, int fd, iree_allocator_t host_allocator,
+ iree_io_stream_t** out_stream);
+
#ifdef __cplusplus
} // extern "C"
#endif // __cplusplus
diff --git a/runtime/src/iree/io/stream.c b/runtime/src/iree/io/stream.c
index 09efdca..5e021d4 100644
--- a/runtime/src/iree/io/stream.c
+++ b/runtime/src/iree/io/stream.c
@@ -6,6 +6,10 @@
#include "iree/io/stream.h"
+#include "iree/io/file_handle.h"
+#include "iree/io/memory_stream.h"
+#include "iree/io/stdio_stream.h"
+
// 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 (32 * 1024)
@@ -298,3 +302,107 @@
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_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;
+ }
+ case IREE_IO_FILE_HANDLE_TYPE_FD: {
+ iree_io_stdio_stream_mode_t stdio_mode = 0;
+ if (iree_all_bits_set(mode, IREE_IO_STREAM_MODE_READABLE)) {
+ stdio_mode |= IREE_IO_STDIO_STREAM_MODE_READ;
+ }
+ if (iree_all_bits_set(mode, IREE_IO_STREAM_MODE_WRITABLE)) {
+ stdio_mode |= IREE_IO_STDIO_STREAM_MODE_WRITE;
+ }
+ status = iree_io_stdio_stream_open_fd(stdio_mode, file_primitive.value.fd,
+ host_allocator, &stream);
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_io_stream_seek(stream, IREE_IO_STREAM_SEEK_SET, file_offset);
+ }
+ 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/stream.h b/runtime/src/iree/io/stream.h
index 5cf0979..3265b93 100644
--- a/runtime/src/iree/io/stream.h
+++ b/runtime/src/iree/io/stream.h
@@ -16,6 +16,8 @@
extern "C" {
#endif // __cplusplus
+typedef struct iree_io_file_handle_t iree_io_file_handle_t;
+
//===----------------------------------------------------------------------===//
// iree_io_stream_t
//===----------------------------------------------------------------------===//
@@ -155,7 +157,7 @@
iree_io_stream_pos_t length);
//===----------------------------------------------------------------------===//
-// Lifetime management utilities
+// iree_io_stream_t utilities
//===----------------------------------------------------------------------===//
typedef void(IREE_API_PTR* iree_io_stream_release_fn_t)(
@@ -177,6 +179,27 @@
return callback;
}
+// TODO(benvanik): remove/rework iree_io_stream_open so that it doesn't pull in
+// any implementations by putting callbacks on the file handle constructors.
+// Today this requires we link the stream implementations we support into the
+// same library and that users wanting to add their own may only open them
+// explicitly. That may end up being sufficient.
+
+// 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);
+
//===----------------------------------------------------------------------===//
// iree_io_stream_t implementation details
//===----------------------------------------------------------------------===//
diff --git a/runtime/src/iree/tooling/BUILD.bazel b/runtime/src/iree/tooling/BUILD.bazel
index 39aed05..6b33185 100644
--- a/runtime/src/iree/tooling/BUILD.bazel
+++ b/runtime/src/iree/tooling/BUILD.bazel
@@ -115,9 +115,7 @@
":numpy_io",
"//runtime/src/iree/base",
"//runtime/src/iree/hal",
- "//runtime/src/iree/io:stdio_stream",
"//runtime/src/iree/io:stream",
- "//runtime/src/iree/io:vec_stream",
"//runtime/src/iree/modules/hal",
"//runtime/src/iree/vm",
],
@@ -130,7 +128,7 @@
":function_io",
"//runtime/src/iree/base",
"//runtime/src/iree/hal",
- "//runtime/src/iree/io:vec_stream",
+ "//runtime/src/iree/io:stream",
"//runtime/src/iree/modules/hal",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
@@ -182,8 +180,7 @@
deps = [
":device_util",
":numpy_io",
- "//runtime/src/iree/io:memory_stream",
- "//runtime/src/iree/io:vec_stream",
+ "//runtime/src/iree/io:stream",
"//runtime/src/iree/testing:gtest",
"//runtime/src/iree/testing:gtest_main",
"//runtime/src/iree/tooling/testdata/npy",
@@ -223,7 +220,7 @@
"//runtime/src/iree/base",
"//runtime/src/iree/base/internal:flags",
"//runtime/src/iree/hal",
- "//runtime/src/iree/io:stdio_stream",
+ "//runtime/src/iree/io:stream",
"//runtime/src/iree/modules/hal:types",
"//runtime/src/iree/vm",
"//runtime/src/iree/vm/bytecode:module",
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index 451b3f8..1c1458f 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -130,9 +130,7 @@
::numpy_io
iree::base
iree::hal
- iree::io::stdio_stream
iree::io::stream
- iree::io::vec_stream
iree::modules::hal
iree::vm
PUBLIC
@@ -147,7 +145,7 @@
::function_io
iree::base
iree::hal
- iree::io::vec_stream
+ iree::io::stream
iree::modules::hal
iree::testing::gtest
iree::testing::gtest_main
@@ -208,8 +206,7 @@
DEPS
::device_util
::numpy_io
- iree::io::memory_stream
- iree::io::vec_stream
+ iree::io::stream
iree::testing::gtest
iree::testing::gtest_main
iree::tooling::testdata::npy
@@ -256,7 +253,7 @@
iree::base
iree::base::internal::flags
iree::hal
- iree::io::stdio_stream
+ iree::io::stream
iree::modules::hal::types
iree::vm
iree::vm::bytecode::module
diff --git a/tools/test/BUILD.bazel b/tools/test/BUILD.bazel
index 9369902..23a64db 100644
--- a/tools/test/BUILD.bazel
+++ b/tools/test/BUILD.bazel
@@ -28,6 +28,7 @@
"executable_sources.mlir",
"iree-benchmark-executable.mlir",
"iree-benchmark-module.mlir",
+ "iree-convert-parameters.txt",
"iree-dump-parameters.txt",
"iree-run-mlir.mlir",
"iree-run-module-expected.mlir",
@@ -62,6 +63,7 @@
"//tools:iree-benchmark-executable",
"//tools:iree-benchmark-module",
"//tools:iree-compile",
+ "//tools:iree-convert-parameters",
"//tools:iree-dump-parameters",
"//tools:iree-opt",
"//tools:iree-run-mlir",
diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt
index f577ec6..9d515d8 100644
--- a/tools/test/CMakeLists.txt
+++ b/tools/test/CMakeLists.txt
@@ -23,6 +23,7 @@
"executable_sources.mlir"
"iree-benchmark-executable.mlir"
"iree-benchmark-module.mlir"
+ "iree-convert-parameters.txt"
"iree-dump-parameters.txt"
"iree-run-mlir.mlir"
"iree-run-module-expected.mlir"
@@ -43,6 +44,7 @@
iree-benchmark-executable
iree-benchmark-module
iree-compile
+ iree-convert-parameters
iree-dump-parameters
iree-opt
iree-run-mlir
diff --git a/tools/test/iree-convert-parameters.txt b/tools/test/iree-convert-parameters.txt
new file mode 100644
index 0000000..338610c
--- /dev/null
+++ b/tools/test/iree-convert-parameters.txt
@@ -0,0 +1,12 @@
+// RUN: (iree-convert-parameters \
+// RUN: --parameters=a=%p/parameters_a.safetensors \
+// RUN: --parameters=b=%p/parameters_b.safetensors \
+// RUN: --output=%t.irpa && \
+// RUN: iree-dump-parameters --parameters=%t.irpa) | \
+// RUN: FileCheck %s
+
+// CHECK: Parameter scope <global> (4 entries, 256 total bytes)
+// CHECK: 448 | 480 | 32 | `a0`
+// CHECK: 512 | 544 | 32 | `a1`
+// CHECK: 576 | 640 | 64 | `b0`
+// CHECK: 640 | 768 | 128 | `b1`
