Handle large file sizes on Windows in file_io utils. (#10815)
Fixes #10807.
diff --git a/runtime/src/iree/base/internal/file_io.c b/runtime/src/iree/base/internal/file_io.c
index 1db3e5d..7bcf8bb 100644
--- a/runtime/src/iree/base/internal/file_io.c
+++ b/runtime/src/iree/base/internal/file_io.c
@@ -11,7 +11,6 @@
#if IREE_FILE_IO_ENABLE
#include <errno.h>
-#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
@@ -32,6 +31,14 @@
// be using this method.
#define IREE_FILE_BASE_ALIGNMENT 4096
+#if defined(IREE_PLATFORM_WINDOWS)
+#define iree_fseek64 _fseeki64
+#define iree_ftell64 _ftelli64
+#else
+#define iree_fseek64 fseek
+#define iree_ftell64 ftell
+#endif // IREE_PLATFORM_WINDOWS
+
iree_status_t iree_file_exists(const char* path) {
IREE_ASSERT_ARGUMENT(path);
IREE_TRACE_ZONE_BEGIN(z0);
@@ -46,6 +53,38 @@
return status;
}
+iree_status_t iree_file_query_length(FILE* file, uint64_t* out_length) {
+ IREE_ASSERT_ARGUMENT(out_length);
+ *out_length = 0;
+ if (!file) return iree_ok_status();
+
+ // Capture original offset so we can return to it.
+ uint64_t origin = iree_ftell64(file);
+
+ // Seek to the end of the file.
+ if (iree_fseek64(file, 0, SEEK_END) == -1) {
+ return iree_make_status(iree_status_code_from_errno(errno), "seek (end)");
+ }
+
+ // Query the position, telling us the total file length in bytes.
+ uint64_t file_length = iree_ftell64(file);
+ if (file_length == -1L) {
+ return iree_make_status(iree_status_code_from_errno(errno), "size query");
+ }
+
+ // Seek back to the file origin.
+ if (iree_fseek64(file, origin, SEEK_SET) == -1) {
+ return iree_make_status(iree_status_code_from_errno(errno), "seek (beg)");
+ }
+
+ *out_length = file_length;
+ return iree_ok_status();
+}
+
+bool iree_file_is_at(FILE* file, uint64_t position) {
+ return iree_ftell64(file) == position;
+}
+
iree_status_t iree_file_contents_allocator_ctl(void* self,
iree_allocator_command_t command,
const void* params,
@@ -84,21 +123,17 @@
static iree_status_t iree_file_read_contents_impl(
FILE* file, iree_allocator_t allocator,
iree_file_contents_t** out_contents) {
- // Seek to the end of the file.
- if (fseek(file, 0, SEEK_END) == -1) {
- return iree_make_status(iree_status_code_from_errno(errno), "seek (end)");
+ // Query total file length so we can preallocate storage.
+ // The file size may be larger than the buffer we can allocate (>2GB file on
+ // 32-bit devices) so we check that here early even though it may have false
+ // negatives (not enough virtual address space to allocate, etc).
+ uint64_t file_length = 0;
+ IREE_RETURN_IF_ERROR(iree_file_query_length(file, &file_length));
+ if (file_length > IREE_HOST_SIZE_MAX) {
+ return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
+ "file length exceeds host address range");
}
-
- // Query the position, telling us the total file length in bytes.
- size_t file_size = ftell(file);
- if (file_size == -1L) {
- return iree_make_status(iree_status_code_from_errno(errno), "size query");
- }
-
- // Seek back to the file start.
- if (fseek(file, 0, SEEK_SET) == -1) {
- return iree_make_status(iree_status_code_from_errno(errno), "seek (beg)");
- }
+ iree_host_size_t file_size = (iree_host_size_t)file_length;
// Compute total size with alignment padding.
// We allocate +1 to force a trailing \0 in case this is used as a cstring.
diff --git a/runtime/src/iree/base/internal/file_io.h b/runtime/src/iree/base/internal/file_io.h
index 3418c62..c5f5f8f 100644
--- a/runtime/src/iree/base/internal/file_io.h
+++ b/runtime/src/iree/base/internal/file_io.h
@@ -7,6 +7,9 @@
#ifndef IREE_BASE_INTERNAL_FILE_IO_H_
#define IREE_BASE_INTERNAL_FILE_IO_H_
+#include <stdint.h>
+#include <stdio.h>
+
#include "iree/base/api.h"
#ifdef __cplusplus
@@ -20,6 +23,12 @@
// Returns IREE_STATUS_NOT_FOUND if the file does not exist.
iree_status_t iree_file_exists(const char* path);
+// Returns the remaining length of |file| in bytes from the current position.
+iree_status_t iree_file_query_length(FILE* file, uint64_t* out_length);
+
+// Returns true if |file| position is at |position|.
+bool iree_file_is_at(FILE* file, uint64_t position);
+
// Loaded file contents.
typedef struct iree_file_contents_t {
iree_allocator_t allocator;
diff --git a/runtime/src/iree/tooling/BUILD b/runtime/src/iree/tooling/BUILD
index cca6909..feff3dc 100644
--- a/runtime/src/iree/tooling/BUILD
+++ b/runtime/src/iree/tooling/BUILD
@@ -142,6 +142,7 @@
":numpy_io",
"//runtime/src/iree/base",
"//runtime/src/iree/base:tracing",
+ "//runtime/src/iree/base/internal:file_io",
"//runtime/src/iree/hal",
"//runtime/src/iree/modules/hal",
"//runtime/src/iree/vm",
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index 8ed8dff..a91c336 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -160,6 +160,7 @@
DEPS
::numpy_io
iree::base
+ iree::base::internal::file_io
iree::base::tracing
iree::hal
iree::modules::hal
diff --git a/runtime/src/iree/tooling/vm_util.c b/runtime/src/iree/tooling/vm_util.c
index 7aa0cb7..32b454b 100644
--- a/runtime/src/iree/tooling/vm_util.c
+++ b/runtime/src/iree/tooling/vm_util.c
@@ -11,34 +11,12 @@
#include <stdio.h>
#include "iree/base/api.h"
+#include "iree/base/internal/file_io.h"
#include "iree/base/tracing.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/module.h"
#include "iree/tooling/numpy_io.h"
-// TODO(benvanik): drop use of stdio and make an iree_io_stream_t.
-#if defined(IREE_PLATFORM_WINDOWS)
-static uint64_t iree_file_query_length(FILE* file) {
- _fseeki64(file, 0, SEEK_END);
- uint64_t file_length = _ftelli64(file);
- _fseeki64(file, 0, SEEK_SET);
- return file_length;
-}
-static bool iree_file_is_eof(FILE* file, uint64_t file_length) {
- return _ftelli64(file) == file_length;
-}
-#else
-static uint64_t iree_file_query_length(FILE* file) {
- fseek(file, 0, SEEK_END);
- uint64_t file_length = ftell(file);
- fseek(file, 0, SEEK_SET);
- return file_length;
-}
-static bool iree_file_is_eof(FILE* file, uint64_t file_length) {
- return ftell(file) == file_length;
-}
-#endif // IREE_PLATFORM_*
-
static iree_status_t iree_allocate_and_copy_cstring_from_view(
iree_allocator_t allocator, iree_string_view_t view, char** cstring) {
IREE_RETURN_IF_ERROR(
@@ -62,15 +40,15 @@
file_path.data);
}
- uint64_t file_length = iree_file_query_length(file);
+ uint64_t file_length = 0;
+ iree_status_t status = iree_file_query_length(file, &file_length);
iree_hal_buffer_params_t buffer_params = {0};
buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
buffer_params.access = IREE_HAL_MEMORY_ACCESS_READ;
buffer_params.type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL;
- iree_status_t status = iree_ok_status();
- while (iree_status_is_ok(status) && !iree_file_is_eof(file, file_length)) {
+ while (iree_status_is_ok(status) && !iree_file_is_at(file, file_length)) {
iree_hal_buffer_view_t* buffer_view = NULL;
status = iree_numpy_npy_load_ndarray(
file, IREE_NUMPY_NPY_LOAD_OPTION_DEFAULT, buffer_params,
