| // 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/gguf/gguf_parser.h" | |
| #include "iree/io/formats/irpa/irpa_builder.h" | |
| #include "iree/io/formats/irpa/irpa_parser.h" | |
| #include "iree/io/formats/safetensors/safetensors_parser.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) { | |
| if (format == "gguf") { | |
| CheckApiStatus( | |
| iree_io_parse_gguf_index(file_handle.raw_ptr(), self.raw_ptr()), | |
| "Could not parse gguf file into index"); | |
| } else if (format == "irpa") { | |
| CheckApiStatus( | |
| iree_io_parse_irpa_index(file_handle.raw_ptr(), self.raw_ptr()), | |
| "Could not parse IREE parameter archive file into index"); | |
| } else if (format == "safetensors") { | |
| CheckApiStatus( | |
| iree_io_parse_safetensors_index(file_handle.raw_ptr(), self.raw_ptr()), | |
| "Could not parse safetensors file into index"); | |
| } else { | |
| throw std::invalid_argument( | |
| "Unrecognized file format. Expected one of: 'gguf', 'irpa', " | |
| "'safetensors'"); | |
| } | |
| } | |
| 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); | |
| } | |
| } // namespace | |
| void SetupIoBindings(py::module_ &m) { | |
| m.def("create_io_parameters_module", &CreateIoParametersModule); | |
| py::class_<FileHandle>(m, "FileHandle") | |
| .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); | |
| py::class_<ParameterProvider>(m, "ParameterProvider"); | |
| 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( | |
| "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 |