| // 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 |