Skip to content

Commit

Permalink
Supporting file descriptors in iree_io_stream_open. (#19665)
Browse files Browse the repository at this point in the history
  • Loading branch information
benvanik authored Jan 10, 2025
1 parent 106371d commit 1d91bec
Show file tree
Hide file tree
Showing 13 changed files with 317 additions and 245 deletions.
52 changes: 14 additions & 38 deletions runtime/src/iree/io/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
Expand All @@ -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",
Expand Down Expand Up @@ -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",
],
Expand All @@ -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",
Expand Down
53 changes: 8 additions & 45 deletions runtime/src/iree/io/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
91 changes: 1 addition & 90 deletions runtime/src/iree/io/file_handle.c
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
23 changes: 0 additions & 23 deletions runtime/src/iree/io/file_handle.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
#include <stdint.h>

#include "iree/base/api.h"
#include "iree/io/stream.h"

#ifdef __cplusplus
extern "C" {
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 1d91bec

Please sign in to comment.