| // Copyright 2023 The IREE Authors |
| // |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| #include "./io.h" |
| |
| #include <iostream> |
| #include <string_view> |
| #include <unordered_map> |
| |
| #include "./buffer_interop.h" |
| #include "./vm.h" |
| #include "iree/base/internal/file_io.h" |
| #include "iree/base/internal/path.h" |
| #include "iree/io/formats/irpa/irpa_builder.h" |
| #include "iree/io/formats/parser_registry.h" |
| #include "iree/io/parameter_index_provider.h" |
| #include "iree/modules/io/parameters/module.h" |
| #include "iree/schemas/parameter_archive.h" |
| |
| namespace iree::python { |
| |
| namespace { |
| |
| VmModule CreateIoParametersModule(VmInstance &instance, py::args providers) { |
| iree_vm_module_t *module = nullptr; |
| std::vector<iree_io_parameter_provider_t *> c_providers; |
| iree_host_size_t size = providers.size(); |
| c_providers.resize(size); |
| for (iree_host_size_t i = 0; i < size; ++i) { |
| ParameterProvider *provider = py::cast<ParameterProvider *>(providers[i]); |
| c_providers[i] = provider->raw_ptr(); |
| } |
| CheckApiStatus(iree_io_parameters_module_create( |
| instance.raw_ptr(), size, c_providers.data(), |
| iree_allocator_system(), &module), |
| "Error creating io_parameters module"); |
| return VmModule::StealFromRawPtr(module); |
| } |
| |
| FileHandle FileHandleWrapMemory(py::object host_buffer, bool readable, |
| bool writable, size_t &out_buffer_size) { |
| struct Retained { |
| Retained(py::object host_buffer) |
| : buffer_request(host_buffer, PyBUF_SIMPLE), |
| host_buffer(std::move(host_buffer)) {} |
| PyBufferRequest buffer_request; |
| py::object host_buffer; |
| }; |
| std::unique_ptr<Retained> outer_retained = |
| std::make_unique<Retained>(std::move(host_buffer)); |
| iree_io_file_access_t access = 0; |
| if (readable) access |= IREE_IO_FILE_ACCESS_READ; |
| if (writable) access |= IREE_IO_FILE_ACCESS_WRITE; |
| iree_io_file_handle_t *created_handle; |
| out_buffer_size = outer_retained->buffer_request.view().len; |
| CheckApiStatus( |
| iree_io_file_handle_wrap_host_allocation( |
| access, |
| iree_byte_span_t{ |
| static_cast<uint8_t *>(outer_retained->buffer_request.view().buf), |
| static_cast<iree_host_size_t>( |
| outer_retained->buffer_request.view().len)}, |
| iree_io_file_handle_release_callback_t{ |
| +[](void *user_data, iree_io_file_handle_primitive_t primitive) { |
| Retained *inner_retained = static_cast<Retained *>(user_data); |
| delete inner_retained; |
| }, |
| (void *)outer_retained.get(), |
| }, |
| iree_allocator_system(), &created_handle), |
| "Could not wrap host memory into a file handle"); |
| outer_retained.release(); |
| return FileHandle::StealFromRawPtr(created_handle); |
| } |
| |
| void ParameterIndexAddFromFileHandle(ParameterIndex &self, std::string &key, |
| FileHandle &file_handle, uint64_t length, |
| uint64_t offset, |
| std::optional<std::string> metadata) { |
| iree_io_parameter_index_entry_t entry; |
| memset(&entry, 0, sizeof(entry)); |
| entry.key = iree_make_string_view(key.data(), key.size()); |
| if (metadata) { |
| entry.metadata.data = reinterpret_cast<const uint8_t *>(metadata->data()); |
| entry.metadata.data_length = metadata->size(); |
| } |
| entry.length = length; |
| entry.type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE; |
| entry.storage.file.handle = file_handle.raw_ptr(); |
| entry.storage.file.offset = offset; |
| CheckApiStatus(iree_io_parameter_index_add(self.raw_ptr(), &entry), |
| "Could not add parameter index entry"); |
| } |
| |
| void ParameterIndexParseFileHandle(ParameterIndex &self, |
| FileHandle &file_handle, |
| std::string &format) { |
| CheckApiStatus(iree_io_parse_file_index( |
| iree_make_string_view(format.data(), format.size()), |
| file_handle.raw_ptr(), self.raw_ptr()), |
| "Could not parse parameter file index"); |
| } |
| |
| void ParameterIndexLoadFile(ParameterIndex &self, std::string &file_path, |
| std::optional<std::string> format, bool readable, |
| bool writable, bool mmap) { |
| // Default format from extension. |
| if (!format) { |
| iree_string_view_t path_ext = iree_file_path_extension( |
| iree_make_string_view(file_path.data(), file_path.size())); |
| format.emplace(path_ext.data, path_ext.size); |
| } |
| |
| // Open file. |
| iree_file_read_flags_t read_flags = IREE_FILE_READ_FLAG_DEFAULT; |
| if (mmap) { |
| read_flags = IREE_FILE_READ_FLAG_MMAP; |
| } else { |
| read_flags = IREE_FILE_READ_FLAG_PRELOAD; |
| } |
| iree_file_contents_t *file_contents = nullptr; |
| CheckApiStatus( |
| iree_file_read_contents(file_path.c_str(), read_flags, |
| iree_allocator_system(), &file_contents), |
| "Error opening parameter file"); |
| iree_io_file_handle_release_callback_t release_callback = { |
| +[](void *user_data, iree_io_file_handle_primitive_t handle_primitive) { |
| iree_file_contents_t *file_contents = (iree_file_contents_t *)user_data; |
| iree_file_contents_free(file_contents); |
| }, |
| file_contents, |
| }; |
| |
| // Wrap contents. |
| iree_io_file_handle_t *raw_file_handle = nullptr; |
| iree_status_t status = iree_io_file_handle_wrap_host_allocation( |
| IREE_IO_FILE_ACCESS_READ, file_contents->buffer, release_callback, |
| iree_allocator_system(), &raw_file_handle); |
| if (!iree_status_is_ok(status)) { |
| iree_file_contents_free(file_contents); |
| CheckApiStatus(status, "Error accessing parameter memory"); |
| } |
| |
| // Parse. |
| FileHandle file_handle = FileHandle::StealFromRawPtr(raw_file_handle); |
| ParameterIndexParseFileHandle(self, file_handle, *format); |
| } |
| |
| // Wraps an index and an entry, extending lifetime of the index. |
| struct ParameterIndexEntryWrapper { |
| ParameterIndexEntryWrapper(ParameterIndex index) : index(std::move(index)) {} |
| |
| ParameterIndex index; |
| const iree_io_parameter_index_entry_t *entry = nullptr; |
| }; |
| |
| } // namespace |
| |
| int FileHandle::HandleBufferProtocol(Py_buffer *view, int flags) { |
| auto primitive = iree_io_file_handle_primitive(raw_ptr()); |
| if (primitive.type != IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION) { |
| PyErr_SetString(PyExc_ValueError, |
| "FileHandle is not based on a host allocation and " |
| "cannot be mapped"); |
| return -1; |
| } |
| if (view == NULL) { |
| PyErr_SetString(PyExc_ValueError, "NULL view in getbuffer"); |
| return -1; |
| } |
| |
| view->buf = primitive.value.host_allocation.data; |
| view->len = primitive.value.host_allocation.data_length; |
| bool is_writable = |
| iree_io_file_handle_access(raw_ptr()) & IREE_IO_FILE_ACCESS_WRITE; |
| view->readonly = !is_writable; |
| view->itemsize = 1; |
| view->format = (char *)"B"; // Byte |
| view->ndim = 1; |
| view->shape = nullptr; |
| view->strides = nullptr; |
| view->suboffsets = nullptr; |
| view->internal = nullptr; |
| return 0; |
| } |
| |
| void SetupIoBindings(py::module_ &m) { |
| m.def("create_io_parameters_module", &CreateIoParametersModule); |
| |
| auto file_handle = py::class_<FileHandle>(m, "FileHandle"); |
| BindBufferProtocol<FileHandle>(file_handle); |
| file_handle |
| .def_static( |
| "wrap_memory", |
| [](py::object host_buffer, bool readable, bool writable) { |
| size_t unused_len; |
| return FileHandleWrapMemory(std::move(host_buffer), readable, |
| writable, unused_len); |
| }, |
| py::arg("host_buffer"), py::arg("readable") = true, |
| py::arg("writable") = false) |
| .def_prop_ro( |
| "is_host_allocation", |
| [](FileHandle &self) { |
| auto primitive = iree_io_file_handle_primitive(self.raw_ptr()); |
| return primitive.type == IREE_IO_FILE_HANDLE_TYPE_HOST_ALLOCATION; |
| }) |
| .def_prop_ro( |
| "host_allocation", |
| [](py::handle self) { |
| return py::steal<py::object>(PyMemoryView_FromObject(self.ptr())); |
| }) |
| .def("__repr__", [](py::handle self_object) { |
| if (py::cast<py::bool_>(self_object.attr("is_host_allocation"))) { |
| return py::str("FileHandle<host_allocation({})>") |
| .format(self_object.attr("host_allocation")); |
| } else { |
| return py::str("<FileHandle unknown>"); |
| } |
| }); |
| |
| py::class_<ParameterProvider>(m, "ParameterProvider"); |
| py::class_<ParameterIndexEntryWrapper>(m, "ParameterIndexEntry") |
| .def_prop_ro("key", |
| [](ParameterIndexEntryWrapper &self) { |
| return py::str(self.entry->key.data, self.entry->key.size); |
| }) |
| .def_prop_ro( |
| "length", |
| [](ParameterIndexEntryWrapper &self) { return self.entry->length; }) |
| .def_prop_ro("metadata", |
| [](ParameterIndexEntryWrapper &self) { |
| return py::bytes((const char *)self.entry->metadata.data, |
| self.entry->metadata.data_length); |
| }) |
| .def_prop_ro("is_file", |
| [](ParameterIndexEntryWrapper &self) { |
| return self.entry->type == |
| IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE; |
| }) |
| .def_prop_ro("is_splat", |
| [](ParameterIndexEntryWrapper &self) { |
| return self.entry->type == |
| IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT; |
| }) |
| .def_prop_ro( |
| "file_storage", |
| [](ParameterIndexEntryWrapper &self) { |
| if (self.entry->type != |
| IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_FILE) { |
| throw std::invalid_argument("Entry is not file storage based"); |
| } |
| return py::make_tuple( |
| FileHandle::BorrowFromRawPtr(self.entry->storage.file.handle), |
| self.entry->storage.file.offset); |
| }) |
| .def_prop_ro("file_view", |
| [](py::handle self_object) { |
| auto file_storage = self_object.attr("file_storage"); |
| py::handle file_handle = file_storage[0]; |
| auto offset = py::cast<iree_host_size_t>(file_storage[1]); |
| auto length = |
| py::cast<iree_host_size_t>(self_object.attr("length")); |
| py::object memview = file_handle.attr("host_allocation"); |
| py::slice slice(offset, offset + length); |
| return memview.attr("__getitem__")(slice); |
| }) |
| .def_prop_ro("splat_pattern", |
| [](ParameterIndexEntryWrapper &self) { |
| if (self.entry->type != |
| IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT) { |
| throw std::invalid_argument("Entry is not splat"); |
| } |
| return py::bytes( |
| (const char *)self.entry->storage.splat.pattern, |
| self.entry->storage.splat.pattern_length); |
| }) |
| .def("__repr__", [](py::handle &self_object) { |
| if (py::cast<py::bool_>(self_object.attr("is_splat"))) { |
| return py::str("<ParameterIndexEntry '{}' splat {}:{}>") |
| .format(self_object.attr("key"), |
| self_object.attr("splat_pattern"), |
| self_object.attr("length")); |
| } else if (py::cast<py::bool_>(self_object.attr("is_file"))) { |
| py::object file_storage = self_object.attr("file_storage"); |
| return py::str("<ParameterIndexEntry '{}' {}:{}:{}") |
| .format(self_object.attr("key"), file_storage[0], file_storage[1], |
| self_object.attr("length")); |
| } else { |
| return py::str("<ParameterIndexEntry unknown>"); |
| } |
| }); |
| py::class_<ParameterIndex>(m, "ParameterIndex") |
| .def("__init__", |
| [](ParameterIndex *new_self) { |
| iree_io_parameter_index_t *created; |
| CheckApiStatus(iree_io_parameter_index_create( |
| iree_allocator_system(), &created), |
| "Could not create IO parameter index"); |
| new (new_self) ParameterIndex(); |
| *new_self = ParameterIndex::StealFromRawPtr(created); |
| }) |
| .def("__len__", |
| [](ParameterIndex &self) { |
| return iree_io_parameter_index_count(self.raw_ptr()); |
| }) |
| .def( |
| "__getitem__", |
| [](ParameterIndex &self, iree_host_size_t i) { |
| ParameterIndexEntryWrapper entry_wrapper(self); |
| CheckApiStatus(iree_io_parameter_index_get(self.raw_ptr(), i, |
| &entry_wrapper.entry), |
| "Could not enumerate parameter index"); |
| return entry_wrapper; |
| }, |
| py::arg("i")) |
| .def("items", |
| [](ParameterIndex &self) { |
| py::list items; |
| for (iree_host_size_t i = 0; |
| i < iree_io_parameter_index_count(self.raw_ptr()); ++i) { |
| ParameterIndexEntryWrapper entry_wrapper(self); |
| CheckApiStatus(iree_io_parameter_index_get(self.raw_ptr(), i, |
| &entry_wrapper.entry), |
| "Could not enumerate parameter index"); |
| py::str key(entry_wrapper.entry->key.data, |
| entry_wrapper.entry->key.size); |
| py::object value = py::cast(std::move(entry_wrapper)); |
| items.append(py::make_tuple(key, value)); |
| } |
| return items; |
| }) |
| .def("__repr__", |
| [](ParameterIndex &self) { |
| iree_string_builder_t b; |
| iree_string_builder_initialize(iree_allocator_system(), &b); |
| iree_status_t status = iree_io_parameter_index_dump( |
| iree_string_view_empty(), self.raw_ptr(), &b); |
| iree_string_view_t sv = iree_string_builder_view(&b); |
| py::str result = py::str(sv.data, sv.size); |
| iree_string_builder_deinitialize(&b); |
| CheckApiStatus(status, "Failed to dump parameter index"); |
| return result; |
| }) |
| .def( |
| "reserve", |
| [](ParameterIndex &self, iree_host_size_t new_capacity) { |
| CheckApiStatus( |
| iree_io_parameter_index_reserve(self.raw_ptr(), new_capacity), |
| "Could not reserve capacity"); |
| }, |
| py::arg("new_capacity")) |
| .def( |
| "add_splat", |
| [](ParameterIndex &self, std::string key, py::object pattern, |
| uint64_t total_length, std::optional<std::string> metadata) { |
| iree_io_parameter_index_entry_t entry; |
| memset(&entry, 0, sizeof(entry)); |
| entry.key = iree_make_string_view(key.data(), key.size()); |
| if (metadata) { |
| entry.metadata.data = |
| reinterpret_cast<const uint8_t *>(metadata->data()); |
| entry.metadata.data_length = metadata->size(); |
| } |
| entry.length = total_length; |
| entry.type = IREE_IO_PARAMETER_INDEX_ENTRY_STORAGE_TYPE_SPLAT; |
| PyBufferRequest pattern_info(pattern, PyBUF_SIMPLE); |
| auto pattern_size = pattern_info.view().len; |
| if (pattern_size > sizeof(entry.storage.splat.pattern)) { |
| throw std::invalid_argument( |
| "pattern must be limited to 16 bytes"); |
| } |
| entry.storage.splat.pattern_length = pattern_size; |
| std::memcpy(entry.storage.splat.pattern, pattern_info.view().buf, |
| pattern_size); |
| CheckApiStatus(iree_io_parameter_index_add(self.raw_ptr(), &entry), |
| "Could not add parameter index entry"); |
| }, |
| py::arg("key"), py::arg("pattern"), py::arg("total_length"), |
| py::arg("metadata") = py::none()) |
| .def("add_from_file_handle", ParameterIndexAddFromFileHandle, |
| py::arg("key"), py::arg("file_handle"), py::arg("length"), |
| py::arg("offset") = 0, py::arg("metadata") = py::none()) |
| .def( |
| "add_buffer", |
| [](ParameterIndex &self, std::string key, py::object buffer, |
| bool readable, bool writable, |
| std::optional<std::string> metadata) { |
| size_t buffer_size; |
| FileHandle file_handle = FileHandleWrapMemory( |
| std::move(buffer), readable, writable, buffer_size); |
| ParameterIndexAddFromFileHandle(self, key, file_handle, buffer_size, |
| /*offset=*/0, std::move(metadata)); |
| }, |
| py::arg("key"), py::arg("buffer"), py::arg("readable") = true, |
| py::arg("writable") = false, py::arg("metadata") = py::none()) |
| .def("load_from_file_handle", ParameterIndexParseFileHandle, |
| py::arg("file_handle"), py::arg("format")) |
| .def("load", ParameterIndexLoadFile, py::arg("file_path"), |
| py::arg("format") = py::none(), py::arg("readable") = true, |
| py::arg("writable") = false, py::arg("mmap") = true) |
| .def( |
| "create_provider", |
| [](ParameterIndex &self, std::string scope, |
| std::optional<iree_host_size_t> max_concurrent_operations) { |
| if (!max_concurrent_operations) { |
| max_concurrent_operations = |
| IREE_IO_PARAMETER_INDEX_PROVIDER_DEFAULT_MAX_CONCURRENT_OPERATIONS; |
| } |
| iree_io_parameter_provider_t *created; |
| CheckApiStatus( |
| iree_io_parameter_index_provider_create( |
| iree_make_string_view(scope.data(), scope.size()), |
| self.raw_ptr(), *max_concurrent_operations, |
| iree_allocator_system(), &created), |
| "Could not create parameter provider from index"); |
| return ParameterProvider::StealFromRawPtr(created); |
| }, |
| py::arg("scope") = std::string(), |
| py::arg("max_concurrent_operations") = py::none()) |
| .def( |
| "create_archive_file", |
| [](ParameterIndex &self, std::string file_path, |
| iree_io_physical_offset_t file_offset, |
| ParameterIndex *explicit_target_index) { |
| // If no target index was given, RAII manage a local target index. |
| iree_io_parameter_index_t *target_index = nullptr; |
| ParameterIndex default_target_index; |
| if (explicit_target_index) { |
| target_index = explicit_target_index->raw_ptr(); |
| } else { |
| iree_io_parameter_index_t *created; |
| CheckApiStatus(iree_io_parameter_index_create( |
| iree_allocator_system(), &created), |
| "Could not create IO parameter index"); |
| default_target_index = ParameterIndex::StealFromRawPtr(created); |
| target_index = default_target_index.raw_ptr(); |
| } |
| |
| // Open the file via callback. |
| struct OpenParams { |
| const char *path; |
| }; |
| OpenParams file_open_user_data{file_path.c_str()}; |
| auto file_open_callback = |
| +[](void *user_data, iree_io_physical_offset_t archive_offset, |
| iree_io_physical_size_t archive_length, |
| iree_io_file_handle_t **out_file_handle) -> iree_status_t { |
| OpenParams *params = static_cast<OpenParams *>(user_data); |
| iree_file_contents_t *file_contents = NULL; |
| IREE_RETURN_IF_ERROR(iree_file_create_mapped( |
| params->path, archive_offset + archive_length, archive_offset, |
| (iree_host_size_t)archive_length, iree_allocator_system(), |
| &file_contents)); |
| iree_io_file_handle_release_callback_t release_callback; |
| memset(&release_callback, 0, sizeof(release_callback)); |
| release_callback.fn = |
| +[](void *user_data, |
| iree_io_file_handle_primitive_t handle_primitive) { |
| iree_file_contents_free( |
| static_cast<iree_file_contents_t *>(user_data)); |
| }; |
| release_callback.user_data = file_contents; |
| iree_status_t status = iree_io_file_handle_wrap_host_allocation( |
| IREE_IO_FILE_ACCESS_WRITE, file_contents->buffer, |
| release_callback, iree_allocator_system(), out_file_handle); |
| if (!iree_status_is_ok(status)) { |
| iree_file_contents_free(file_contents); |
| } |
| return status; |
| }; |
| |
| // Write the archive. |
| CheckApiStatus(iree_io_build_parameter_archive( |
| self.raw_ptr(), target_index, |
| iree_io_parameter_archive_file_open_callback_t{ |
| file_open_callback, |
| &file_open_user_data, |
| }, |
| file_offset, iree_allocator_system()), |
| "Error building parameter archive"); |
| |
| // Return the target index. |
| return ParameterIndex::BorrowFromRawPtr(target_index); |
| }, |
| py::arg("file_path"), py::arg("file_offset") = 0, |
| py::arg("target_index") = nullptr); |
| } |
| |
| } // namespace iree::python |
