Pulling misc iree_io_* fixes/cleanups from iree-org#15983. (iree-org#…

…17914)

This is all of the tweaks made during iree-org#15983 as part of adding the new
IO stream dialect. That PR will be revived at some point in the future
without these changes needing to be bundled in.

runtime/src/iree/base/internal/file_io.c

@@ -495,7 +495,7 @@ iree_status_t iree_file_write_contents(const char* path,
 
   iree_status_t status = iree_ok_status();
   if (content.data_length > 0) {
-    int ret = fwrite((char*)content.data, content.data_length, 1, file);
+    size_t ret = fwrite((char*)content.data, content.data_length, 1, file);
     if (ret != 1) {
       status = iree_make_status(IREE_STATUS_DATA_LOSS,
                                 "unable to write file contents of %" PRIhsz

runtime/src/iree/io/file_handle.c

@@ -155,7 +155,7 @@ IREE_API_EXPORT iree_status_t iree_io_stream_open(
             file_offset, file_primitive.value.host_allocation.data_length);
         break;
       }
-      iree_io_memory_stream_release_callback_t release_callback = {
+      iree_io_stream_release_callback_t release_callback = {
           .fn = iree_io_memory_stream_file_release,
           .user_data = file_handle,
       };

runtime/src/iree/io/memory_stream.c

@@ -101,7 +101,7 @@ static iree_status_t iree_io_stream_apply_fixed_seek(
 typedef struct iree_io_memory_stream_t {
   iree_io_stream_t base;
   iree_allocator_t host_allocator;
-  iree_io_memory_stream_release_callback_t release_callback;
+  iree_io_stream_release_callback_t release_callback;
   iree_io_stream_pos_t offset;
   iree_io_stream_pos_t length;
   uint8_t* contents;
@@ -116,7 +116,7 @@ static iree_io_memory_stream_t* iree_io_memory_stream_cast(
 
 IREE_API_EXPORT iree_status_t iree_io_memory_stream_wrap(
     iree_io_stream_mode_t mode, iree_byte_span_t contents,
-    iree_io_memory_stream_release_callback_t release_callback,
+    iree_io_stream_release_callback_t release_callback,
     iree_allocator_t host_allocator, iree_io_stream_t** out_stream) {
   IREE_ASSERT_ARGUMENT(out_stream);
   *out_stream = NULL;

runtime/src/iree/io/memory_stream.h

@@ -18,31 +18,12 @@ extern "C" {
 // iree_io_memory_stream_t
 //===----------------------------------------------------------------------===//
 
-typedef void(IREE_API_PTR* iree_io_memory_stream_release_fn_t)(
-    void* user_data, iree_io_stream_t* stream);
-
-// A callback issued when a memory stream is released.
-typedef struct {
-  // Callback function pointer.
-  iree_io_memory_stream_release_fn_t fn;
-  // User data passed to the callback function. Unowned.
-  void* user_data;
-} iree_io_memory_stream_release_callback_t;
-
-// Returns a no-op file release callback that implies that no cleanup is
-// required.
-static inline iree_io_memory_stream_release_callback_t
-iree_io_memory_stream_release_callback_null(void) {
-  iree_io_memory_stream_release_callback_t callback = {NULL, NULL};
-  return callback;
-}
-
 // Wraps a fixed-size host memory allocation |contents| in a stream.
 // |release_callback| can be used to receive a callback when the stream is
 // destroyed and the reference to the contents is no longer required.
 IREE_API_EXPORT iree_status_t iree_io_memory_stream_wrap(
     iree_io_stream_mode_t mode, iree_byte_span_t contents,
-    iree_io_memory_stream_release_callback_t release_callback,
+    iree_io_stream_release_callback_t release_callback,
     iree_allocator_t host_allocator, iree_io_stream_t** out_stream);
 
 #ifdef __cplusplus

runtime/src/iree/io/memory_stream_test.cc

@@ -28,7 +28,7 @@ TEST(MemoryStreamTest, Wrap) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   EXPECT_EQ(iree_io_stream_mode(stream), IREE_IO_STREAM_MODE_READABLE);
@@ -44,7 +44,7 @@ TEST(MemoryStreamTest, WrapEmpty) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, 0),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   EXPECT_EQ(iree_io_stream_mode(stream), IREE_IO_STREAM_MODE_READABLE);
@@ -57,7 +57,7 @@ TEST(MemoryStreamTest, WrapEmpty) {
 
 TEST(MemoryStreamTest, WrapReleaseCallback) {
   int callback_count = 0;
-  iree_io_memory_stream_release_callback_t release_callback = {
+  iree_io_stream_release_callback_t release_callback = {
       +[](void* user_data, iree_io_stream_t* stream) {
         int* callback_count = (int*)user_data;
         ++(*callback_count);
@@ -83,7 +83,7 @@ TEST(MemoryStreamTest, SeekSet) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -126,7 +126,7 @@ TEST(MemoryStreamTest, SeekFromCurrent) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -188,7 +188,7 @@ TEST(MemoryStreamTest, SeekFromEnd) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -231,7 +231,7 @@ TEST(MemoryStreamTest, SeekToAlignment) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -283,7 +283,7 @@ TEST(MemoryStreamTest, ReadUpTo) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -344,7 +344,7 @@ TEST(MemoryStreamTest, ReadExact) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   // Streams start at origin 0.
@@ -403,7 +403,7 @@ TEST(MemoryStreamTest, Write) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_WRITABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   const uint8_t write_buffer[8] = {0, 1, 2, 3, 4, 5, 6, 7};
@@ -460,7 +460,7 @@ TEST(MemoryStreamTest, Fill) {
   iree_io_stream_t* stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_WRITABLE, iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   uint8_t pattern[] = {0x80, 0x90, 0xA0, 0xB0, 0xC0, 0xD0, 0xE0, 0xF0};
@@ -534,7 +534,7 @@ TEST(MemoryStreamTest, MapRead) {
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE | IREE_IO_STREAM_MODE_MAPPABLE,
       iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   iree_const_byte_span_t span = iree_const_byte_span_empty();
@@ -565,7 +565,7 @@ TEST(MemoryStreamTest, MapWrite) {
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_WRITABLE | IREE_IO_STREAM_MODE_MAPPABLE,
       iree_make_byte_span(data, sizeof(data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &stream));
 
   iree_byte_span_t span = iree_byte_span_empty();
@@ -597,15 +597,15 @@ TEST(MemoryStreamTest, Copy) {
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE,
       iree_make_byte_span(source_data, sizeof(source_data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &source_stream));
 
   uint8_t target_data[5] = {0xDD};
   iree_io_stream_t* target_stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_WRITABLE,
       iree_make_byte_span(target_data, sizeof(target_data)),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &target_stream));
 
   // Bounds checks length.
@@ -670,15 +670,15 @@ TEST(MemoryStreamTest, CopyLarge) {
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_READABLE,
       iree_make_byte_span(source_data.data(), source_data.size()),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &source_stream));
 
   std::vector<uint8_t> target_data(1 * 1024 * 1024);
   iree_io_stream_t* target_stream = NULL;
   IREE_ASSERT_OK(iree_io_memory_stream_wrap(
       IREE_IO_STREAM_MODE_WRITABLE,
       iree_make_byte_span(target_data.data(), target_data.size()),
-      iree_io_memory_stream_release_callback_null(), iree_allocator_system(),
+      iree_io_stream_release_callback_null(), iree_allocator_system(),
       &target_stream));
 
   // Copy an interior subrange.

runtime/src/iree/io/stream.c

@@ -67,6 +67,14 @@ static iree_status_t iree_io_stream_validate_mode(
 // iree_io_stream_t
 //===----------------------------------------------------------------------===//
 
+IREE_API_EXPORT void iree_io_stream_destroy(iree_io_stream_t* stream) {
+  if (IREE_LIKELY(stream)) {
+    IREE_TRACE_ZONE_BEGIN(z0);
+    stream->vtable->destroy(stream);
+    IREE_TRACE_ZONE_END(z0);
+  }
+}
+
 IREE_API_EXPORT void iree_io_stream_retain(iree_io_stream_t* stream) {
   if (IREE_LIKELY(stream)) {
     iree_atomic_ref_count_inc(&stream->ref_count);
@@ -76,9 +84,7 @@ IREE_API_EXPORT void iree_io_stream_retain(iree_io_stream_t* stream) {
 IREE_API_EXPORT void iree_io_stream_release(iree_io_stream_t* stream) {
   if (IREE_LIKELY(stream) &&
       iree_atomic_ref_count_dec(&stream->ref_count) == 1) {
-    IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree_io_stream_destroy");
-    stream->vtable->destroy(stream);
-    IREE_TRACE_ZONE_END(z0);
+    iree_io_stream_destroy(stream);
   }
 }
 

runtime/src/iree/io/stream.h

@@ -154,6 +154,29 @@ IREE_API_EXPORT iree_status_t iree_io_stream_copy(
     iree_io_stream_t* source_stream, iree_io_stream_t* target_stream,
     iree_io_stream_pos_t length);
 
+//===----------------------------------------------------------------------===//
+// Lifetime management utilities
+//===----------------------------------------------------------------------===//
+
+typedef void(IREE_API_PTR* iree_io_stream_release_fn_t)(
+    void* user_data, iree_io_stream_t* stream);
+
+// A callback issued when a memory stream is released.
+typedef struct {
+  // Callback function pointer.
+  iree_io_stream_release_fn_t fn;
+  // User data passed to the callback function. Unowned.
+  void* user_data;
+} iree_io_stream_release_callback_t;
+
+// Returns a no-op file release callback that implies that no cleanup is
+// required.
+static inline iree_io_stream_release_callback_t
+iree_io_stream_release_callback_null(void) {
+  iree_io_stream_release_callback_t callback = {NULL, NULL};
+  return callback;
+}
+
 //===----------------------------------------------------------------------===//
 // iree_io_stream_t implementation details
 //===----------------------------------------------------------------------===//
@@ -191,6 +214,8 @@ struct iree_io_stream_t {
   iree_io_stream_mode_t mode;
 };
 
+IREE_API_EXPORT void iree_io_stream_destroy(iree_io_stream_t* stream);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif  // __cplusplus

runtime/src/iree/modules/vmvx/CMakeLists.txt

@@ -17,8 +17,8 @@ iree_cc_library(
   TEXTUAL_HDRS
     "exports.inl"
   SRCS
-  "elementwise.c"
-  "module.c"
+    "elementwise.c"
+    "module.c"
   DEFINES
     "IREE_HAVE_VMVX_MODULE"
   DEPS

runtime/src/iree/tooling/numpy_io_test.cc

@@ -48,8 +48,8 @@ class NumpyIOTest : public ::testing::Test {
       IREE_CHECK_OK(iree_io_memory_stream_wrap(
           IREE_IO_STREAM_MODE_READABLE | IREE_IO_STREAM_MODE_SEEKABLE,
           iree_make_byte_span((void*)file_toc[i].data, file_toc[i].size),
-          iree_io_memory_stream_release_callback_null(),
-          iree_allocator_system(), &stream));
+          iree_io_stream_release_callback_null(), iree_allocator_system(),
+          &stream));
       return StreamPtr(stream, iree_io_stream_release);
     }
     return StreamPtr{nullptr, iree_io_stream_release};