From 1d91bec41b10ca83753c670f74c697e54ee3d979 Mon Sep 17 00:00:00 2001 From: Ben Vanik Date: Fri, 10 Jan 2025 09:44:17 -0800 Subject: [PATCH] Supporting file descriptors in iree_io_stream_open. (#19665) --- runtime/src/iree/io/BUILD.bazel | 52 ++------ runtime/src/iree/io/CMakeLists.txt | 53 ++------ runtime/src/iree/io/file_handle.c | 91 +------------ runtime/src/iree/io/file_handle.h | 23 ---- runtime/src/iree/io/stdio_stream.c | 169 +++++++++++++++++++----- runtime/src/iree/io/stdio_stream.h | 7 + runtime/src/iree/io/stream.c | 108 +++++++++++++++ runtime/src/iree/io/stream.h | 25 +++- runtime/src/iree/tooling/BUILD.bazel | 9 +- runtime/src/iree/tooling/CMakeLists.txt | 9 +- tools/test/BUILD.bazel | 2 + tools/test/CMakeLists.txt | 2 + tools/test/iree-convert-parameters.txt | 12 ++ 13 files changed, 317 insertions(+), 245 deletions(-) create mode 100644 tools/test/iree-convert-parameters.txt diff --git a/runtime/src/iree/io/BUILD.bazel b/runtime/src/iree/io/BUILD.bazel index 7d4c0e1005a0..892989f79ade 100644 --- a/runtime/src/iree/io/BUILD.bazel +++ b/runtime/src/iree/io/BUILD.bazel @@ -17,19 +17,6 @@ iree_runtime_cc_library( 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 @@ iree_runtime_cc_test( 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", @@ -92,33 +79,22 @@ iree_runtime_cc_library( ], ) -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", + ], + hdrs = [ + "memory_stream.h", + "stdio_stream.h", + "stream.h", + "vec_stream.h", ], -) - -iree_runtime_cc_library( - name = "vec_stream", - srcs = ["vec_stream.c"], - hdrs = ["vec_stream.h"], deps = [ - ":stream", + ":file_handle", "//runtime/src/iree/base", "//runtime/src/iree/base/internal", ], @@ -128,7 +104,7 @@ iree_runtime_cc_test( 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 cf4fd9f88acd..b6ad802445a6 100644 --- a/runtime/src/iree/io/CMakeLists.txt +++ b/runtime/src/iree/io/CMakeLists.txt @@ -18,22 +18,6 @@ iree_cc_library( 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 @@ iree_cc_test( SRCS "memory_stream_test.cc" DEPS - ::memory_stream + ::stream iree::base iree::testing::gtest iree::testing::gtest_main @@ -109,42 +93,21 @@ iree_cc_library( PUBLIC ) -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 @@ iree_cc_test( 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 9a6727bcf31c..22d9d804dd2b 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 @@ static iree_status_t iree_io_platform_map_file_view( 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_io_file_mapping_contents_rw(iree_io_file_mapping_t* mapping) { 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 908f1e7ae81a..ba2f81b48e56 100644 --- a/runtime/src/iree/io/file_handle.h +++ b/runtime/src/iree/io/file_handle.h @@ -10,7 +10,6 @@ #include #include "iree/base/api.h" -#include "iree/io/stream.h" #ifdef __cplusplus extern "C" { @@ -308,28 +307,6 @@ iree_io_file_mapping_contents_ro(const iree_io_file_mapping_t* mapping); 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 b190585e5b1b..b26362c1e646 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 + #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 @@ IREE_API_EXPORT iree_status_t iree_io_stdio_stream_wrap( } #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 @@ IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open( 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_API_EXPORT iree_status_t iree_io_stdio_stream_open( 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 @@ IREE_API_EXPORT iree_status_t iree_io_stdio_stream_open( "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 @@ static iree_status_t iree_io_stdio_stream_read( // 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 @@ static iree_status_t iree_io_stdio_stream_read( 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 ddb277332322..aef08eff2072 100644 --- a/runtime/src/iree/io/stdio_stream.h +++ b/runtime/src/iree/io/stdio_stream.h @@ -51,6 +51,13 @@ 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); +// 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 09efdca5c21a..5e021d46f76c 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_API_EXPORT iree_status_t iree_io_stream_copy( 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 5cf09794eca1..3265b930ab12 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_API_EXPORT iree_status_t iree_io_stream_copy( 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 @@ iree_io_stream_release_callback_null(void) { 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 39aed05ae7ba..6b33185f130c 100644 --- a/runtime/src/iree/tooling/BUILD.bazel +++ b/runtime/src/iree/tooling/BUILD.bazel @@ -115,9 +115,7 @@ iree_runtime_cc_library( ":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 @@ iree_runtime_cc_test( ":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 @@ iree_runtime_cc_test( 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 @@ iree_runtime_cc_library( "//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 451b3f839bb4..1c1458ffc234 100644 --- a/runtime/src/iree/tooling/CMakeLists.txt +++ b/runtime/src/iree/tooling/CMakeLists.txt @@ -130,9 +130,7 @@ iree_cc_library( ::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 @@ iree_cc_test( ::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 @@ iree_cc_test( 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_cc_library( 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 936990276a3b..23a64dbe83d7 100644 --- a/tools/test/BUILD.bazel +++ b/tools/test/BUILD.bazel @@ -28,6 +28,7 @@ iree_lit_test_suite( "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 @@ iree_lit_test_suite( "//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 f577ec643775..9d515d8d94ff 100644 --- a/tools/test/CMakeLists.txt +++ b/tools/test/CMakeLists.txt @@ -23,6 +23,7 @@ iree_lit_test_suite( "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_lit_test_suite( 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 000000000000..338610c37b70 --- /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 (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`