| #include "neural_runtime.hpp" |
|
|
| #include <algorithm> |
| #include <array> |
| #include <cstdio> |
| #include <cstdlib> |
| #include <filesystem> |
| #include <limits> |
| #include <memory> |
| #include <string> |
| #include <thread> |
| #include <utility> |
| #include <vector> |
|
|
| #if defined(MOTIONBRICKS_HAVE_GGML) |
| #include <ggml-cpu.h> |
| #include <gguf.h> |
| #if defined(MOTIONBRICKS_HAVE_VULKAN) |
| #include <ggml-vulkan.h> |
| #endif |
| #endif |
|
|
| namespace motionbricks::detail { |
|
|
| #if defined(MOTIONBRICKS_HAVE_GGML) |
| namespace { |
|
|
| struct ggml_context_deleter { |
| void operator()(ggml_context * value) const noexcept { ggml_free(value); } |
| }; |
| struct gguf_context_deleter { |
| void operator()(gguf_context * value) const noexcept { gguf_free(value); } |
| }; |
| struct backend_deleter { |
| void operator()(ggml_backend * value) const noexcept { ggml_backend_free(value); } |
| }; |
| struct buffer_deleter { |
| void operator()(ggml_backend_buffer * value) const noexcept { ggml_backend_buffer_free(value); } |
| }; |
|
|
| using context_ptr = std::unique_ptr<ggml_context, ggml_context_deleter>; |
| using gguf_ptr = std::unique_ptr<gguf_context, gguf_context_deleter>; |
| using backend_ptr = std::unique_ptr<ggml_backend, backend_deleter>; |
| using buffer_ptr = std::unique_ptr<ggml_backend_buffer, buffer_deleter>; |
|
|
| bool copy_gguf_data(const std::filesystem::path & path, ggml_context * tensors, |
| const gguf_context * metadata, std::string & reason) { |
| std::unique_ptr<std::FILE, decltype(&std::fclose)> file( |
| std::fopen(path.string().c_str(), "rb"), &std::fclose); |
| if (!file) { |
| reason = "cannot open GGUF weights: " + path.string(); |
| return false; |
| } |
| std::vector<std::byte> staging(4U * 1024U * 1024U); |
| const auto count = gguf_get_n_tensors(metadata); |
| for (std::int64_t index = 0; index < count; ++index) { |
| const char * name = gguf_get_tensor_name(metadata, index); |
| auto * tensor = ggml_get_tensor(tensors, name); |
| if (tensor == nullptr) { |
| reason = std::string("missing allocated tensor: ") + name; |
| return false; |
| } |
| const auto offset = gguf_get_data_offset(metadata) + gguf_get_tensor_offset(metadata, index); |
| if (offset > static_cast<std::size_t>(std::numeric_limits<long>::max()) || |
| std::fseek(file.get(), static_cast<long>(offset), SEEK_SET) != 0) { |
| reason = std::string("cannot seek to tensor data: ") + name; |
| return false; |
| } |
| const auto bytes = ggml_nbytes(tensor); |
| for (std::size_t position = 0; position < bytes; position += staging.size()) { |
| const auto amount = std::min(staging.size(), bytes - position); |
| if (std::fread(staging.data(), 1, amount, file.get()) != amount) { |
| reason = std::string("truncated tensor data: ") + name; |
| return false; |
| } |
| ggml_backend_tensor_set(tensor, staging.data(), position, amount); |
| } |
| } |
| return true; |
| } |
|
|
| } |
|
|
| class neural_runtime { |
| public: |
| struct component { |
| std::string name; |
| context_ptr context; |
| buffer_ptr buffer; |
| }; |
|
|
| backend_ptr backend; |
| std::vector<component> components; |
| }; |
|
|
| namespace { |
|
|
| mb_status open_component(neural_runtime & runtime, const std::filesystem::path & path, |
| std::string name, std::string & reason) { |
| ggml_context * raw_context = nullptr; |
| const gguf_init_params params{true, &raw_context}; |
| gguf_ptr metadata(gguf_init_from_file(path.string().c_str(), params)); |
| context_ptr context(raw_context); |
| if (!metadata || !context) { |
| reason = "cannot initialize GGUF weights: " + path.string(); |
| return MB_INVALID_FORMAT; |
| } |
| buffer_ptr buffer(ggml_backend_alloc_ctx_tensors(context.get(), runtime.backend.get())); |
| if (!buffer) { |
| reason = "cannot allocate backend weights for " + path.string(); |
| return MB_OUT_OF_MEMORY; |
| } |
| if (!copy_gguf_data(path, context.get(), metadata.get(), reason)) return MB_IO_ERROR; |
| runtime.components.push_back({std::move(name), std::move(context), std::move(buffer)}); |
| return MB_OK; |
| } |
|
|
| } |
|
|
| mb_status create_neural_runtime(const std::filesystem::path & bundle, |
| mb_device device, std::uint32_t threads, |
| const std::string & backend_directory, |
| std::shared_ptr<neural_runtime> & output, |
| std::string & reason) { |
| output.reset(); |
| if (!backend_directory.empty()) ggml_backend_load_all_from_path(backend_directory.c_str()); |
| auto runtime = std::make_shared<neural_runtime>(); |
| const auto selected = device == MB_DEVICE_AUTO ? MB_DEVICE_CPU : device; |
| if (selected == MB_DEVICE_CPU) { |
| runtime->backend.reset(ggml_backend_cpu_init()); |
| if (runtime->backend) { |
| const auto hardware = std::max(1U, std::thread::hardware_concurrency()); |
| const auto count = threads == 0U ? hardware : threads; |
| ggml_backend_cpu_set_n_threads(runtime->backend.get(), static_cast<int>(count)); |
| } |
| } else if (selected == MB_DEVICE_VULKAN) { |
| #if defined(MOTIONBRICKS_HAVE_VULKAN) |
| |
| |
| |
| |
| #if defined(_WIN32) |
| _putenv_s("GGML_VK_DISABLE_F16", "1"); |
| _putenv_s("GGML_VK_DISABLE_COOPMAT", "1"); |
| _putenv_s("GGML_VK_DISABLE_COOPMAT2", "1"); |
| #else |
| setenv("GGML_VK_DISABLE_F16", "1", 0); |
| setenv("GGML_VK_DISABLE_COOPMAT", "1", 0); |
| setenv("GGML_VK_DISABLE_COOPMAT2", "1", 0); |
| #endif |
| runtime->backend.reset(ggml_backend_vk_init(0)); |
| #else |
| reason = "this build has no Vulkan backend"; |
| return MB_BACKEND_UNAVAILABLE; |
| #endif |
| } |
| if (!runtime->backend) { |
| reason = selected == MB_DEVICE_VULKAN ? "cannot initialize Vulkan backend" |
| : "cannot initialize CPU backend"; |
| return MB_BACKEND_UNAVAILABLE; |
| } |
| constexpr std::array files{ |
| std::pair{"pose", "pose.gguf"}, std::pair{"root", "root.gguf"}, |
| std::pair{"vq-decoder", "vq-decoder.gguf"}, std::pair{"support", "support.gguf"}, |
| }; |
| for (const auto & [name, filename] : files) { |
| const auto status = open_component(*runtime, bundle / filename, name, reason); |
| if (status != MB_OK) return status; |
| } |
| output = std::move(runtime); |
| return MB_OK; |
| } |
|
|
| ggml_backend_t neural_backend(const neural_runtime & runtime) noexcept { |
| return runtime.backend.get(); |
| } |
|
|
| ggml_tensor * neural_weight(const neural_runtime & runtime, |
| std::string_view component, std::string_view name) noexcept { |
| const auto found = std::find_if(runtime.components.begin(), runtime.components.end(), |
| [&](const neural_runtime::component & item) { return item.name == component; }); |
| if (found == runtime.components.end()) return nullptr; |
| std::string owned(name); |
| if (auto * direct = ggml_get_tensor(found->context.get(), owned.c_str())) return direct; |
| if (owned.size() >= 64U) { |
| const std::string needle = "self_attn."; |
| if (const auto position = owned.find(needle); position != std::string::npos) |
| owned.replace(position, needle.size(), "attn."); |
| } |
| return ggml_get_tensor(found->context.get(), owned.c_str()); |
| } |
|
|
| bool neural_copy_f32(const neural_runtime & runtime, |
| std::string_view component, std::string_view name, |
| std::vector<float> & output, std::string & reason) { |
| auto * tensor = neural_weight(runtime, component, name); |
| if (tensor == nullptr) { |
| reason = "missing neural tensor: " + std::string(component) + ":" + std::string(name); |
| return false; |
| } |
| if (tensor->type != GGML_TYPE_F32) { |
| reason = "neural tensor is not F32: " + std::string(component) + ":" + std::string(name); |
| return false; |
| } |
| output.resize(static_cast<std::size_t>(ggml_nelements(tensor))); |
| ggml_backend_tensor_get(tensor, output.data(), 0, output.size() * sizeof(float)); |
| return true; |
| } |
|
|
| #else |
|
|
| class neural_runtime {}; |
|
|
| mb_status create_neural_runtime(const std::filesystem::path &, mb_device, std::uint32_t, |
| const std::string &, std::shared_ptr<neural_runtime> & output, |
| std::string & reason) { |
| output.reset(); |
| reason = "this build has no GGML support"; |
| return MB_BACKEND_UNAVAILABLE; |
| } |
|
|
| #endif |
|
|
| } |
|
|
