From fb85f0137483984f05a964ce7152a54b6e753cbd Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Fri, 30 Oct 2026 19:59:54 +0000 Subject: [PATCH 01/18] ggml-hrx: manage weight representations by execution demand Materialize imported weights in selected layouts, retain demand per source region, and retire obsolete representations after GPU users complete. Keep other immutable GGUF weights on reusable native storage. Assisted-by: Codex --- ggml/include/ggml-hrx.h | 10 + ggml/src/ggml-hrx/README.md | 12 + ggml/src/ggml-hrx/backend-buffer-binding.cpp | 148 +- ggml/src/ggml-hrx/backend-buffer-binding.h | 23 + ggml/src/ggml-hrx/backend-context.h | 20 +- .../dispatch/command-program-bindings.cpp | 60 +- .../dispatch/command-program-bindings.h | 22 + .../dispatch/command-program-resolver.cpp | 2 +- .../common/dispatch-gated-mul-mat.cpp | 28 +- .../common/dispatch-mul-mat.cpp | 37 +- .../dispatch-registry.cpp | 31 +- .../dispatch_registration/dispatch-registry.h | 3 + .../llm/dispatch-attention-qkv.cpp | 3 + ggml/src/ggml-hrx/ggml-hrx.cpp | 212 +- ggml/src/ggml-hrx/graph/value-map.cpp | 3 + ggml/src/ggml-hrx/graph/value-map.h | 4 + .../runtime/command-program-executor.cpp | 268 +- .../runtime/command-program-executor.h | 41 + ggml/src/ggml-hrx/runtime/graph-executor.cpp | 437 ++- ggml/src/ggml-hrx/runtime/graph-executor.h | 1 + .../ggml-hrx/runtime/graph-program-cache.cpp | 305 +- .../ggml-hrx/runtime/graph-program-cache.h | 46 +- .../ggml-hrx/runtime/host-buffer-registry.cpp | 29 +- .../ggml-hrx/runtime/host-buffer-registry.h | 23 +- ggml/src/ggml-hrx/runtime/host-memory.cpp | 224 +- ggml/src/ggml-hrx/runtime/host-memory.h | 32 +- .../prepared-command-program-cache.cpp | 73 +- .../runtime/prepared-command-program-cache.h | 7 +- .../runtime/weight-submission-tracker.h | 77 + tests/hrx-backend-test.cpp | 469 ++- tests/hrx-buffer-test.cpp | 349 ++- tests/hrx-ops-test.cpp | 2517 ++++++++++++++++- tests/test-llama-archs.cpp | 107 +- 33 files changed, 5337 insertions(+), 286 deletions(-) create mode 100644 ggml/src/ggml-hrx/runtime/weight-submission-tracker.h diff --git a/ggml/include/ggml-hrx.h b/ggml/include/ggml-hrx.h index d40c095b75bc..999160950b3b 100644 --- a/ggml/include/ggml-hrx.h +++ b/ggml/include/ggml-hrx.h @@ -11,12 +11,22 @@ struct ggml_backend_hrx_cache_stats { uint64_t graph_program_hits; uint64_t prepared_program_builds; uint64_t prepared_program_hits; + uint64_t weight_materializations; + uint64_t weight_reuses; + uint64_t weight_uploaded_bytes; + uint64_t weight_live_allocations; + uint64_t weight_live_bytes; + uint64_t weight_peak_bytes; + uint64_t host_uploads; + uint64_t host_upload_bytes; + uint64_t weight_duplicate_warnings; }; GGML_BACKEND_API ggml_backend_t ggml_backend_hrx_init(size_t device); GGML_BACKEND_API bool ggml_backend_is_hrx(ggml_backend_t backend); GGML_BACKEND_API int ggml_backend_hrx_get_device_count(void); GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device); +GGML_BACKEND_API bool ggml_backend_hrx_buffer_is_imported(ggml_backend_buffer_t buffer); GGML_BACKEND_API bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, struct ggml_backend_hrx_cache_stats * stats); GGML_BACKEND_API ggml_backend_reg_t ggml_backend_hrx_reg(void); diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index 47a530eaf4ef..cfa97b1331d9 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -1,5 +1,17 @@ # HRX Backend +## Weight residency diagnostics + +Weight residency supports standard GGUF inference with immutable model weights. Alternative layouts require an HRX `buffer_from_host_ptr` import that does not alias a known HRX host allocation. Other weights remain native-only: host weights reuse a cached native copy, and directly bindable device weights reuse their existing allocation. Source availability for native caching does not grant permission to select a transformed layout. + +Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx: weight_stats` line when each HRX backend is released. The `-v` flag enables the backend log callback in `llama-bench`. The line reports successful weight materializations, cache acquisition reuse, uploaded weight bytes, live allocation count and bytes, peak live bytes, and all host transfer uploads. `ggml_backend_hrx_get_cache_stats` exposes the same weight counters while a backend is running. Prepared program hits do not count as weight cache acquisitions, so unchanged replay can reuse a weight without increasing the `reused` counter. + +Live weight bytes count physical allocations retained by the cache, a prepared program, or a recorded graph. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. + +Each successfully admitted HRX submission updates the retained layout set only for the weight regions it reads. A later fragment using other weights, or no weights, leaves those regions resident. When a region is used with a different set, completed obsolete layouts are released while unfinished submissions keep their bindings until completion. A coherent native read also replaces that region's layout set, even when it needs no device allocation. Repeating warm disjoint fragments reuses their allocations; explicit cache clearing and backend teardown require materialization again. + +Keep borrowed weight bytes unchanged and alive while HRX can use them. A foreign CPU host buffer and its storage must stay alive for the lifetime of the HRX backend that caches it. Its identity depends on its wrapper and base addresses; destroying and recreating a buffer at those addresses while that backend lives can reuse old residency. Training and in-place weight updates are outside this path. Ordinary mutable inputs and outputs retain their host staging, writeback and synchronization behavior. + ## `hrx-system` backend compatibility `hrx-minimum-commit.txt` records the minimum required commit in ROCm/hrx-system. It contains one full Git commit SHA. The backend's C++ APIs, Loom syntax, and numerical behavior can all depend on that revision. When `GGML_HRX=ON` and `HRX_SOURCE_DIR` is set, CMake requires the HRX checkout's HEAD to equal or descend from the recorded commit. diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp index bab5f12b1af7..5e21c06e1dc8 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.cpp +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -3,6 +3,39 @@ #include "ggml-backend-impl.h" #include "ggml.h" +#include +#include + +namespace { + +std::mutex foreign_host_generations_mutex; +// GGML does not notify HRX when a foreign backend destroys a buffer. The source lifetime contract holds while an +// HRX backend is alive, and the table is retired with the last such backend. +std::unordered_map foreign_host_generations; +size_t active_hrx_backends = 0; + +uint64_t foreign_host_generation(uint64_t identity) { + std::lock_guard lock(foreign_host_generations_mutex); + const auto found = foreign_host_generations.find(identity); + return found == foreign_host_generations.end() ? 1 : found->second; +} + +} // namespace + +void ggml_backend_hrx_foreign_host_backend_started() { + std::lock_guard lock(foreign_host_generations_mutex); + ++active_hrx_backends; +} + +void ggml_backend_hrx_foreign_host_backend_stopped() { + std::lock_guard lock(foreign_host_generations_mutex); + GGML_ASSERT(active_hrx_backends > 0); + --active_hrx_backends; + if (active_hrx_backends == 0) { + foreign_host_generations.clear(); + } +} + ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer) { return static_cast(buffer->context); } @@ -27,7 +60,13 @@ bool ggml_backend_hrx_tensor_binding(const ggml_tensor * tensor, } auto * context = ggml_backend_hrx_buffer_context_from_buffer(buffer); const size_t offset = ggml_backend_hrx_tensor_offset(context, tensor); - if (context->buffer == nullptr || offset > buffer->size || ggml_nbytes(tensor) > buffer->size - offset) { + const bool has_device_buffer = context->buffer != nullptr; + const bool is_imported = context->imported; + const bool has_storage = has_device_buffer || is_imported; + const bool offset_fits = offset <= buffer->size; + const bool tensor_size_fits = offset_fits ? ggml_nbytes(tensor) <= buffer->size - offset : false; + const bool valid_binding = has_storage && offset_fits && tensor_size_fits; + if (!valid_binding) { return false; } *out_context = context; @@ -65,22 +104,117 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx binding.length = ggml_nbytes(tensor); binding.identity = buffer_address ^ (base_address + 0x9e3779b97f4a7c15ull + (buffer_address << 6) + (buffer_address >> 2)); - binding.generation = 1; + binding.generation = foreign_host_generation(binding.identity); binding.capacity = capacity; binding.weight = ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + // Standard GGUF weights are immutable, so their native copy can use the shared weight cache. + // Layout eligibility remains a separate import-only decision. + binding.source_host_data = base; + binding.source_offset = host_offset; + binding.source_capacity = capacity; return true; } ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; - const bool directly_bindable = !ggml_backend_buffer_is_host(buffer) || context->direct_host_binding; + const bool is_imported = context->imported; + const bool is_host_buffer = ggml_backend_buffer_is_host(buffer); + const bool ordinary_device_buffer = !is_imported && !is_host_buffer; + const bool directly_bindable = ordinary_device_buffer || context->direct_host_binding; // Coherent HRX host allocations are directly device-addressable. Represent them with an HRX buffer handle so // command-program preparation bypasses host materialization. Noncoherent host allocations remain host data. binding.buffer = directly_bindable ? context->buffer : nullptr; - binding.host_data = directly_bindable ? nullptr : context->base; - binding.offset = offset; + const bool shared_host_region = context->retained_base != nullptr; + binding.host_data = directly_bindable ? nullptr : + (shared_host_region ? context->retained_base : context->base); + binding.offset = shared_host_region ? context->physical_offset + offset : offset; binding.length = ggml_nbytes(tensor); binding.identity = context->identity; - binding.generation = context->generation; - binding.capacity = buffer != nullptr ? buffer->size : 0; + binding.generation = ggml_backend_hrx_buffer_generation(context); + binding.capacity = shared_host_region ? context->physical_capacity : buffer->size; binding.weight = buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + const bool retained_host_source = is_imported || is_host_buffer; + binding.source_host_data = retained_host_source ? + (shared_host_region ? context->retained_base : context->base) : nullptr; + binding.source_offset = binding.offset; + binding.source_capacity = binding.capacity; return true; } + +bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region) { + ggml_backend_hrx_buffer_context * context = nullptr; + size_t offset = 0; + const bool hrx_bound = ggml_backend_hrx_tensor_binding(tensor, &context, &offset); + if (hrx_bound) { + ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + const bool retained = context->imported || ggml_backend_buffer_is_host(buffer); + if (!retained) { + return false; + } + region.identity = context->identity; + region.generation = ggml_backend_hrx_buffer_generation(context); + region.offset = context->retained_base != nullptr ? context->physical_offset + offset : offset; + } else { + ggml::hrx::ValueBufferBinding binding; + const bool resolved = ggml_backend_hrx_resolve_value_buffer(tensor, binding); + const bool has_host_source = resolved && binding.host_data != nullptr; + if (!has_host_source) { + return false; + } + region.identity = binding.identity; + region.generation = binding.generation; + region.offset = binding.offset; + } + region.length = ggml_nbytes(tensor); + region.type = tensor->type; + for (size_t i = 0; i < GGML_MAX_DIMS; ++i) { + region.ne[i] = tensor->ne[i]; + region.nb[i] = tensor->nb[i]; + } + return true; +} + +void ggml_backend_hrx_mark_foreign_host_region_written(const ggml_tensor * tensor) { + ggml_backend_hrx_imported_region region = {}; + if (!ggml_backend_hrx_get_retained_region(tensor, region)) { + return; + } + std::lock_guard lock(foreign_host_generations_mutex); + uint64_t & generation = foreign_host_generations[region.identity]; + generation = generation == 0 ? 2 : generation + 1; +} + +bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region) { + ggml_backend_hrx_buffer_context * context = nullptr; + size_t offset = 0; + const bool bound = ggml_backend_hrx_tensor_binding(tensor, &context, &offset); + return bound && context->imported && ggml_backend_hrx_get_retained_region(tensor, region); +} + +bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return false; + } + const ggml_tensor * root = tensor->view_src != nullptr ? tensor->view_src : tensor; + ggml_backend_buffer_t buffer = root->buffer; + if (buffer == nullptr) { + return false; + } + if (buffer->iface.get_base == ggml_backend_hrx_buffer_base) { + const auto * context = ggml_backend_hrx_buffer_context_from_buffer(buffer); + const bool borrowed_import = context->imported; + const bool known_host_alias = context->retained_base != nullptr; + const bool repackable_source = borrowed_import && !known_host_alias; + return repackable_source; + } + return false; +} + +bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return true; + } + const ggml_tensor * root = tensor->view_src != nullptr ? tensor->view_src : tensor; + ggml_backend_buffer_t buffer = root->buffer; + const bool is_weight = buffer != nullptr && + ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; + return !is_weight || ggml_backend_hrx_weight_has_repackable_source(tensor); +} diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.h b/ggml/src/ggml-hrx/backend-buffer-binding.h index 59b925873f63..bd110f4abe13 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.h +++ b/ggml/src/ggml-hrx/backend-buffer-binding.h @@ -4,10 +4,23 @@ #include "ggml-backend.h" #include "graph/value-map.h" +#include #include struct ggml_tensor; +// The buffer identity and byte range name the source independently of a resident HRX allocation. +// GGML's tensor metadata supplies the logical shape, including views into a shared import. +struct ggml_backend_hrx_imported_region { + uint64_t identity; + uint64_t generation; + size_t offset; + size_t length; + ggml_type type; + std::array ne; + std::array nb; +}; + ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer); size_t ggml_backend_hrx_tensor_offset(const ggml_backend_hrx_buffer_context * context, const ggml_tensor * tensor); void * ggml_backend_hrx_buffer_base(ggml_backend_buffer_t buffer); @@ -16,3 +29,13 @@ bool ggml_backend_hrx_tensor_binding(const ggml_tensor * tensor, ggml_backend_hrx_buffer_context ** out_context, size_t * out_offset); bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx::ValueBufferBinding & binding); +bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region); +// Only borrowed imports outside HRX-owned host storage may use alternative weight layouts. +bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor); +bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor); +bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region); +void ggml_backend_hrx_mark_imported_region_written(const ggml_tensor * tensor); +void ggml_backend_hrx_mark_retained_region_written(const ggml_tensor * tensor); +void ggml_backend_hrx_mark_foreign_host_region_written(const ggml_tensor * tensor); +void ggml_backend_hrx_foreign_host_backend_started(); +void ggml_backend_hrx_foreign_host_backend_stopped(); diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index 3d510b310b59..1e9e211f2994 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -8,6 +8,7 @@ #include "runtime/kernel-executable-cache.h" #include "runtime/prepared-command-program-cache.h" #include "runtime/transient-arena.h" +#include "runtime/weight-submission-tracker.h" #include #include @@ -32,10 +33,26 @@ struct ggml_backend_hrx_buffer_context { hrx_buffer_t buffer; uint8_t * base; uint64_t identity; - uint64_t generation; + std::atomic generation; bool direct_host_binding; + bool imported; + size_t physical_offset = 0; + size_t physical_capacity = 0; + uint8_t * retained_base = nullptr; + std::shared_ptr> shared_generation; }; +inline uint64_t ggml_backend_hrx_buffer_generation(const ggml_backend_hrx_buffer_context * context) { + return context->shared_generation ? context->shared_generation->load() : context->generation.load(); +} + +inline void ggml_backend_hrx_set_buffer_generation(ggml_backend_hrx_buffer_context * context, uint64_t generation) { + context->generation = generation; + if (context->shared_generation) { + *context->shared_generation = generation; + } +} + struct ggml_backend_hrx_device_context { hrx_device_t device = nullptr; std::string name; @@ -63,6 +80,7 @@ struct ggml_backend_hrx_context { ggml::hrx::TransientArena transient_arena; ggml::hrx::HostTransferManager host_transfers; ggml::hrx::HostWeightCache host_weights; + ggml::hrx::WeightSubmissionTracker weight_submissions; ggml::hrx::GraphReplayStreamState graph_replay_state; std::string name; }; diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 741a41c9c17f..191a8dc8b558 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -13,6 +13,22 @@ static void mix_hash(uint64_t & hash, uint64_t value) { } // namespace +CommandProgramBinding make_command_program_binding( + ValueId value, const ValueBufferBinding & buffer, ggml_type type, + const std::array & ne, const std::array & nb, + bool empty_value) { + CommandProgramBinding binding = { value, buffer.buffer, buffer.offset, buffer.length, buffer.identity, + buffer.generation, buffer.capacity, buffer.host_data, buffer.weight, + empty_value }; + binding.source_host_data = buffer.source_host_data; + binding.source_offset = buffer.source_offset; + binding.source_capacity = buffer.source_capacity; + binding.source_type = type; + binding.source_ne = ne; + binding.source_nb = nb; + return binding; +} + CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & values) { std::vector bindings; CommandProgramBindings result; @@ -27,9 +43,9 @@ CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & v result.status.log("external value %d is not bound", id.value); continue; } - bindings.push_back({ value->id, buffer->buffer, buffer->offset, buffer->length, buffer->identity, - buffer->generation, buffer->capacity, buffer->host_data, buffer->weight, - value->byte_count == 0 }); + CommandProgramBinding binding = make_command_program_binding( + value->id, *buffer, value->type, value->ne, value->nb, value->byte_count == 0); + bindings.push_back(binding); } return from_bindings(std::move(bindings), result.status); } @@ -59,6 +75,23 @@ const CommandProgramBinding * CommandProgramBindings::find(ValueId value) const return nullptr; } +const CommandProgramBinding * CommandProgramBindings::find(const CommandBinding & request) const { + for (const CommandProgramBinding & binding : bindings_) { + const bool same_value = binding.value == request.value; + const bool same_layout = binding.layout == request.layout; + const bool native_layout = request.layout == kNativeWeightLayout; + const bool same_transformation = native_layout || + (binding.layout_source_type == request.source_type && binding.input_size == request.input_size && + binding.output_size == request.output_size && binding.source_length == request.source_length && + binding.length == request.length); + const bool matches_request = same_value && same_layout && same_transformation; + if (matches_request) { + return &binding; + } + } + return nullptr; +} + CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBindings & bindings) { uint64_t hash = UINT64_C(1469598103934665603); mix_hash(hash, UINT64_C(0x6872782d62696e64)); @@ -66,12 +99,20 @@ CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBin mix_hash(hash, static_cast(static_cast(binding.value.value))); mix_hash(hash, binding.host_data != nullptr ? 1 : 0); mix_hash(hash, binding.identity); - mix_hash(hash, binding.generation); + mix_hash(hash, binding.tracks_content_generation() ? binding.generation : 0); mix_hash(hash, static_cast(binding.capacity)); mix_hash(hash, static_cast(binding.offset)); mix_hash(hash, static_cast(binding.length)); mix_hash(hash, binding.weight ? 1 : 0); mix_hash(hash, binding.empty_value ? 1 : 0); + mix_hash(hash, binding.source_host_data != nullptr ? 1 : 0); + mix_hash(hash, static_cast(binding.source_offset)); + mix_hash(hash, static_cast(binding.source_capacity)); + mix_hash(hash, static_cast(binding.source_type)); + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + mix_hash(hash, static_cast(binding.source_ne[dimension])); + mix_hash(hash, static_cast(binding.source_nb[dimension])); + } } mix_hash(hash, static_cast(bindings.bindings().size())); return { hash }; @@ -82,9 +123,16 @@ CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const Com out << "hrx-bindings-v1"; for (const CommandProgramBinding & binding : bindings.bindings()) { out << "|value=" << binding.value.value << "|kind=" << (binding.host_data != nullptr ? "host" : "device") - << "|identity=" << binding.identity << "|generation=" << binding.generation + << "|identity=" << binding.identity << "|generation=" + << (binding.tracks_content_generation() ? binding.generation : 0) << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length - << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0); + << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0) + << "|source_host=" << (binding.source_host_data != nullptr ? 1 : 0) + << "|source_offset=" << binding.source_offset << "|source_capacity=" << binding.source_capacity + << "|source_type=" << static_cast(binding.source_type); + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + out << "|ne=" << binding.source_ne[dimension] << "|nb=" << binding.source_nb[dimension]; + } } return { out.str() }; } diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index d43cc5e8f3ee..5356eae4771c 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -1,10 +1,12 @@ #pragma once #include "graph/value-map.h" +#include "command-program.h" #include "status.h" #include #include +#include #include #include @@ -23,10 +25,29 @@ struct CommandProgramBinding { void * host_data = nullptr; bool weight = false; bool empty_value = false; + // Retained host bytes are separate from the address used to bind a command. + const void * source_host_data = nullptr; + size_t source_offset = 0; + size_t source_capacity = 0; + ggml_type source_type = GGML_TYPE_COUNT; + std::array source_ne = {}; + std::array source_nb = {}; + // Materialized bindings are selected by the requesting command, not just by ValueId. + std::string layout = kNativeWeightLayout; + ggml_type layout_source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; + size_t source_length = 0; bool requires_materialization() const { return host_data != nullptr; } + // Mutable host staging refreshes at launch; retained weight content is prepared ahead of execution. + bool tracks_content_generation() const { return weight && source_host_data != nullptr; } }; +CommandProgramBinding make_command_program_binding( + ValueId value, const ValueBufferBinding & buffer, ggml_type type, + const std::array & ne, const std::array & nb, bool empty_value); + struct CommandProgramBindingsFingerprint { std::string value; }; @@ -42,6 +63,7 @@ class CommandProgramBindings { const Status & errors = {}); const CommandProgramBinding * find(ValueId value) const; + const CommandProgramBinding * find(const CommandBinding & request) const; const std::vector & bindings() const { return bindings_; } diff --git a/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp b/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp index 18805390c903..dfc41d6a676d 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp @@ -23,7 +23,7 @@ static Status resolve_command_binding(const Command & comman switch (binding.origin) { case CommandBindingOrigin::GraphValue: { - const CommandProgramBinding * concrete = bindings.find(binding.value); + const CommandProgramBinding * concrete = bindings.find(binding); if (concrete == nullptr) { status.log("%s %s is not bound", command_context.c_str(), binding_context.c_str()); return status; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp index c0016ba9f7d4..cfeb31398a78 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -723,10 +723,16 @@ static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLU return false; } - const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 && - match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 && - match.gate_format == CommonMulMatWeightFormat::Q4K && - match.up_format == CommonMulMatWeightFormat::Q4K; + const bool has_few_tokens = match.token_count <= 5; + const bool has_aligned_output = match.output_size % 64 == 0; + const bool has_direct_gate = match.gate_weight->alias_source.value < 0; + const bool has_direct_up = match.up_weight->alias_source.value < 0; + const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K; + const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K; + const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight); + const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight); + const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up && + has_q4_gate && has_q4_up && can_transform_gate && can_transform_up; const bool use_q8 = pack_q4; const bool use_direct_dot = match.token_count <= 5 && match.input_size % 256 == 0; return !use_q8 && !use_direct_dot && match.token_count >= 2; @@ -958,10 +964,16 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, } } - const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 && - match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 && - match.gate_format == CommonMulMatWeightFormat::Q4K && - match.up_format == CommonMulMatWeightFormat::Q4K; + const bool has_few_tokens = match.token_count <= 5; + const bool has_aligned_output = match.output_size % 64 == 0; + const bool has_direct_gate = match.gate_weight->alias_source.value < 0; + const bool has_direct_up = match.up_weight->alias_source.value < 0; + const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K; + const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K; + const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight); + const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight); + const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up && + has_q4_gate && has_q4_up && can_transform_gate && can_transform_up; const bool use_q8 = pack_q4; const CommonActivationPublicationDemand q8_demand = use_q8 ? select_swiglu_q8_publication(context, match) : CommonActivationPublicationDemand{}; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp index 047caa6acec3..4e785ac0fb79 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -1034,8 +1034,11 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c } } - const bool use_f16 = match.weight->type == GGML_TYPE_Q4_K && match.weight->alias_source.value < 0 && - match.output_size >= match.input_size / 4; + const bool is_q4_weight = match.weight->type == GGML_TYPE_Q4_K; + const bool has_direct_weight = match.weight->alias_source.value < 0; + const bool has_wide_output = match.output_size >= match.input_size / 4; + const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); + const bool use_f16 = is_q4_weight && has_direct_weight && has_wide_output && can_transform_weight; DispatchBinding activation; const bool packed_input = use_f16 && @@ -1070,7 +1073,8 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c dispatch.kernel.compile_parameters.emplace( use_f16 ? "ggml.workload.token_capacity" : "ggml.mul_mat_q8_1_x4.token_capacity", common_to_config_value(match.token_count)); - const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && match.weight->alias_source.value < 0; + const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && + has_direct_weight && can_transform_weight; if (use_f16) { dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.f16_input_layout", packed_input ? "1" : "0"); } else { @@ -1310,8 +1314,13 @@ static bool build_mul_mat_dispatch( dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_accumulation", "0"); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_unary_op", std::to_string(unary_kind_config_value(match.output_unary_op))); - const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && - match.output_size % 64 == 0 && match.weight->alias_source.value < 0; + const bool has_few_tokens = match.token_count <= 5; + const bool has_skinny_reduction = common_mul_mat_has_skinny_reduction_shape(match.input_size); + const bool has_aligned_output = match.output_size % 64 == 0; + const bool has_direct_weight = match.weight->alias_source.value < 0; + const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); + const bool pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && + has_direct_weight && can_transform_weight; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; bool pack_q4 = false; @@ -1412,7 +1421,14 @@ static bool prepare_vector_activation_binding(const DispatchMatchContext & c weight_layout = nullptr; uses_q8_activation = false; - if (token_count != 1 || input_size % 256 != 0 || weight.alias_source.value >= 0 || output_size % 64 != 0) { + const bool has_single_token = token_count == 1; + const bool has_aligned_input = input_size % 256 == 0; + const bool has_direct_weight = weight.alias_source.value < 0; + const bool has_aligned_output = output_size % 64 == 0; + const bool can_transform_weight = dispatch_weight_layout_eligible(weight); + const bool can_prepare_packed_vector = has_single_token && has_aligned_input && has_direct_weight && + has_aligned_output && can_transform_weight; + if (!can_prepare_packed_vector) { return true; } @@ -1806,8 +1822,13 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.output_size", common_to_config_value(match.output_size)); - const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && - match.output_size % 64 == 0 && match.weight->alias_source.value < 0; + const bool has_few_tokens = match.token_count <= 5; + const bool has_skinny_reduction = common_mul_mat_has_skinny_reduction_shape(match.input_size); + const bool has_aligned_output = match.output_size % 64 == 0; + const bool has_direct_weight = match.weight->alias_source.value < 0; + const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); + const bool pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && + has_direct_weight && can_transform_weight; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; bool pack_q4 = false; diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp index 193f3dac977e..7555e200fa60 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -1,5 +1,7 @@ #include "dispatch-registry.h" +#include "backend-buffer-binding.h" + #include "common/dispatch-common.h" #include "llm/dispatch-attention-qkv.h" #include "llm/dispatch-gated-delta-net.h" @@ -38,6 +40,27 @@ static bool qwen_dispatch_disabled_from_environment() { std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0; } +static bool dispatch_has_eligible_weight_layouts(const DispatchMatchContext & context, const DispatchMatch & match) { + auto eligible = [&](const std::vector & dispatches) { + for (const Dispatch & dispatch : dispatches) { + for (const DispatchBinding & binding : dispatch.bindings) { + if (binding.layout == kNativeWeightLayout) { + continue; + } + const Value * value = context.graph.values().find(binding.value); + const bool has_graph_value = value != nullptr; + const bool layout_is_eligible = !has_graph_value || dispatch_weight_layout_eligible(*value); + const bool has_ineligible_weight = has_graph_value && !layout_is_eligible; + if (has_ineligible_weight) { + return false; + } + } + } + return true; + }; + return eligible(match.initialization_dispatches) && eligible(match.dispatches); +} + static DispatchRegistry build_registry(bool include_qwen) { DispatchRegistryBuilder builder; register_common_dispatches(builder); @@ -52,6 +75,10 @@ static DispatchRegistry build_registry(bool include_qwen) { } // namespace +bool dispatch_weight_layout_eligible(const Value & value) { + return ggml_backend_hrx_weight_layout_eligible(value.tensor); +} + bool DispatchRegistry::match(const DispatchMatchContext & context, DispatchMatch & match) const { return this->match(context, match, nullptr); } @@ -70,7 +97,9 @@ bool DispatchRegistry::match(const DispatchMatchContext & context, registrations_by_root_[static_cast(context.root_node->op)].ordered; for (const DispatchRegistration & registration : registrations) { DispatchMatch candidate; - if (registration.matcher != nullptr && registration.matcher(context, candidate)) { + const bool matched = registration.matcher != nullptr && registration.matcher(context, candidate); + const bool eligible = matched && dispatch_has_eligible_weight_layouts(context, candidate); + if (eligible) { if (diagnostics != nullptr) { diagnostics->attempts.push_back({ registration.name != nullptr ? registration.name : "", diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h index f491890c1bd3..2fe76c617983 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h @@ -48,6 +48,9 @@ struct DispatchMatchContext { const DispatchRegistry * registry = nullptr; }; +// Non-native weight bindings require bytes that remain available after graph import. +bool dispatch_weight_layout_eligible(const Value & value); + enum DispatchActivationInputFormat : uint32_t { DispatchActivationInputNone = 0, DispatchActivationInputF16Row = 1u << 0, diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp index 98d90b2957d3..4e07e2933ad3 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp @@ -312,6 +312,9 @@ static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & co } static bool can_use_packed_lowtoken_attention_matmul(const AttentionMatMulMatch & root) { + if (!dispatch_weight_layout_eligible(*root.weight)) { + return false; + } if (root.token_count < 1 || root.token_count > 5 || root.output_size % 64 != 0 || root.weight->alias_source.value >= 0) { return false; diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 71b4f9977ae6..cbc4a15be8d8 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -180,7 +180,10 @@ static hrx_status_t submit_copy_buffer(hrx_stream_t stream, void * user_data) { static void buffer_free(ggml_backend_buffer_t buffer) { auto * context = buffer_context(buffer); - if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { + const bool owned_buffer = !context->imported; + const bool host_backed = context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + const bool registered_host_buffer = owned_buffer && host_backed; + if (registered_host_buffer) { context->device->host_buffers.remove(context->buffer); } if (context->buffer != nullptr) { @@ -202,6 +205,7 @@ static void buffer_memset(ggml_backend_buffer_t buffer, GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { std::memset(context->base + destination_offset, value, size); + ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } FillBufferArgs args{ context->buffer, destination_offset, size, value }; @@ -222,7 +226,8 @@ static void buffer_set(ggml_backend_buffer_t buffer, const size_t destination_offset = tensor_offset(context, tensor) + offset; GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { - std::memcpy(context->base + destination_offset, data, size); + std::memmove(context->base + destination_offset, data, size); + ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } if (!HRX_CHECK(hrx_synchronous_h2d(context->device->device, data, context->buffer, destination_offset, size))) { @@ -267,8 +272,33 @@ static bool buffer_copy(ggml_backend_buffer_t buffer, const ggml_tensor * source destination_offset > buffer->size || size > buffer->size - destination_offset) { return false; } - CopyBufferArgs args{ source_context->buffer, source_offset, destination_context->buffer, destination_offset, size }; - return buffer_submit_and_wait(destination_context->device, submit_copy_buffer, &args); + const bool source_imported = source_context->imported; + const bool destination_imported = destination_context->imported; + const bool both_imported = source_imported && destination_imported; + if (both_imported) { + std::memmove(destination_context->base + destination_offset, source_context->base + source_offset, size); + } else if (source_imported) { + if (!HRX_CHECK(hrx_synchronous_h2d(destination_context->device->device, source_context->base + source_offset, + destination_context->buffer, destination_offset, size))) { + return false; + } + } else if (destination_imported) { + if (!HRX_CHECK(hrx_synchronous_d2h(source_context->device->device, source_context->buffer, source_offset, + destination_context->base + destination_offset, size))) { + return false; + } + } else { + CopyBufferArgs args{ source_context->buffer, source_offset, destination_context->buffer, destination_offset, + size }; + if (!buffer_submit_and_wait(destination_context->device, submit_copy_buffer, &args)) { + return false; + } + } + const bool destination_host = destination_context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + if (destination_host) { + ggml_backend_hrx_set_buffer_generation(destination_context, g_allocation_generation.fetch_add(1)); + } + return true; } static void buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { @@ -278,6 +308,7 @@ static void buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { auto * context = buffer_context(buffer); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { std::memset(context->base, value, buffer->size); + ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } FillBufferArgs args{ context->buffer, 0, buffer->size, value }; @@ -332,7 +363,8 @@ static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_ } const uint64_t generation = g_allocation_generation.fetch_add(1); auto * context = new (std::nothrow) ggml_backend_hrx_buffer_context{ - type_context->device, allocation, base, generation, generation, direct_host_binding, + type_context->device, allocation, base, generation, generation, direct_host_binding, false, + 0, 0, nullptr, {}, }; if (context == nullptr) { if (allocation != nullptr) { @@ -341,7 +373,11 @@ static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_ return nullptr; } if (host_visible && allocation != nullptr) { - type_context->device->host_buffers.add(allocation, base, size); + context->retained_base = base; + context->physical_capacity = size; + context->shared_generation = std::make_shared>(generation); + type_context->device->host_buffers.add(allocation, base, size, direct_host_binding, + context->identity, context->shared_generation); } return ggml_backend_buffer_init(buft, buffer_i, context, size); } @@ -368,6 +404,15 @@ static void backend_synchronize(ggml_backend_t backend); static void backend_free(ggml_backend_t backend) { auto * context = static_cast(backend->context); backend_synchronize(backend); + if (environment_flag_enabled("GGML_HRX_LOG_WEIGHT_STATS")) { + const ggml::hrx::HostWeightCacheStats weights = context->host_weights.stats(); + const ggml::hrx::HostTransferStats transfers = context->host_transfers.stats(); + GGML_LOG_INFO("ggml_hrx: weight_stats materializations=%" PRIu64 " reused=%" PRIu64 + " uploaded_bytes=%zu live_allocations=%zu live_bytes=%zu peak_bytes=%zu" + " host_uploads=%" PRIu64 " host_upload_bytes=%zu\n", + weights.misses, weights.hits, weights.materialized_bytes, weights.allocation_count, + weights.resident_bytes, weights.peak_resident_bytes, transfers.uploads, transfers.upload_bytes); + } context->prepared_programs.clear(); context->graph_programs.clear(); context->kernel_executables.clear(); @@ -377,6 +422,7 @@ static void backend_free(ggml_backend_t backend) { hrx_stream_release(context->stream); delete context; delete backend; + ggml_backend_hrx_foreign_host_backend_stopped(); } static bool synchronous_upload_fallback(ggml_backend_hrx_context * backend, @@ -436,6 +482,13 @@ static void backend_set_tensor_async(ggml_backend_t backend, GGML_LOG_ERROR("%s: invalid HRX tensor upload\n", __func__); return; } + const bool retained_host = context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + if (retained_host) { + backend_synchronize(backend); + buffer_set(tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer, + tensor, data, offset, size); + return; + } ggml::hrx::HostBufferRef source = backend_context->device->host_buffers.find(data, size); if (source.valid()) { // Registered host buffers can participate directly in the stream command buffer. @@ -459,6 +512,13 @@ static void backend_get_tensor_async(ggml_backend_t backend, GGML_LOG_ERROR("%s: invalid HRX tensor download\n", __func__); return; } + const bool retained_host = context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + if (retained_host) { + backend_synchronize(backend); + buffer_get(tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer, + tensor, data, offset, size); + return; + } ggml::hrx::HostBufferRef destination = backend_context->device->host_buffers.find(data, size); if (destination.valid()) { // Registered host buffers can participate directly in the stream command buffer. @@ -473,7 +533,6 @@ static bool backend_copy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * source, ggml_tensor * destination) { - GGML_UNUSED(backend_src); auto * destination_backend = static_cast(backend_dst->context); ggml_backend_hrx_buffer_context * destination_context = nullptr; size_t destination_offset = 0; @@ -487,11 +546,34 @@ static bool backend_copy_tensor_async(ggml_backend_t backend_src, if (source_context->device != destination_context->device) { return false; } + const bool source_host = source_context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + const bool destination_host = destination_context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + const bool touches_retained_host = source_host || destination_host; + if (touches_retained_host) { + ggml_backend_synchronize(backend_src); + if (backend_dst != backend_src) { + ggml_backend_synchronize(backend_dst); + } + ggml_backend_buffer_t destination_buffer = + destination->view_src != nullptr ? destination->view_src->buffer : destination->buffer; + return buffer_copy(destination_buffer, source, destination); + } return HRX_CHECK(hrx_stream_copy_buffer(destination_backend->stream, source_context->buffer, source_offset, destination_context->buffer, destination_offset, size)); } ggml_backend_buffer_t source_buffer = source->view_src != nullptr ? source->view_src->buffer : source->buffer; if (source_buffer != nullptr && ggml_backend_buffer_is_host(source_buffer)) { + const bool destination_host = destination_context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); + if (destination_host) { + ggml_backend_synchronize(backend_src); + if (backend_dst != backend_src) { + ggml_backend_synchronize(backend_dst); + } + ggml_backend_buffer_t destination_buffer = + destination->view_src != nullptr ? destination->view_src->buffer : destination->buffer; + buffer_set(destination_buffer, destination, source->data, 0, size); + return true; + } return synchronous_upload_fallback(destination_backend, source->data, destination_context->buffer, destination_offset, size); } @@ -502,6 +584,14 @@ static void backend_synchronize(ggml_backend_t backend) { auto * context = static_cast(backend->context); if (HRX_CHECK(hrx_stream_synchronize(context->stream))) { context->graph_replay_state.mark_stream_synchronized(); + const bool obsolete_completed = context->weight_submissions.has_obsolete_pending(); + context->weight_submissions.mark_stream_synchronized(); + if (obsolete_completed) { + const auto & current_demand = context->weight_submissions.current_demand(); + context->graph_programs.invalidate_obsolete_weights(current_demand); + context->prepared_programs.invalidate_obsolete_weights(current_demand); + context->host_weights.evict_except(current_demand, {}); + } } } @@ -562,7 +652,7 @@ static void device_props(ggml_backend_dev_t device, ggml_backend_dev_props * pro device_memory(device, &props->memory_free, &props->memory_total); props->type = GGML_BACKEND_DEVICE_TYPE_GPU; props->device_id = nullptr; - props->caps = { true, true, false, false }; + props->caps = { true, true, true, false }; } static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parameters) { @@ -584,6 +674,8 @@ static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parame if (backend == nullptr) { delete context; hrx_stream_release(stream); + } else { + ggml_backend_hrx_foreign_host_backend_started(); } return backend; } @@ -596,6 +688,46 @@ static ggml_backend_buffer_type_t device_host_buffer_type(ggml_backend_dev_t dev return &device_context(device)->host_buft; } +static ggml_backend_buffer_t device_buffer_from_host_ptr(ggml_backend_dev_t device, + void * ptr, + size_t size, + size_t max_tensor_size) { + GGML_UNUSED(max_tensor_size); + const bool has_pointer = ptr != nullptr; + const bool has_size = size != 0; + const bool valid_region = has_pointer && has_size; + if (!valid_region) { + return nullptr; + } + const uint64_t generation = g_allocation_generation.fetch_add(1); + auto * context = new (std::nothrow) ggml_backend_hrx_buffer_context{ + device_context(device), nullptr, static_cast(ptr), generation, generation, false, true, + 0, 0, nullptr, {}, + }; + if (context == nullptr) { + return nullptr; + } + ggml::hrx::HostBufferRef known_host = device_context(device)->host_buffers.find(ptr, size); + const bool known_host_region = known_host.valid(); + const bool coherent_host_region = known_host.direct_bindable(); + const bool known_direct_host = known_host_region && coherent_host_region; + if (known_host_region) { + context->identity = known_host.identity(); + context->shared_generation = known_host.generation(); + context->retained_base = static_cast(known_host.base()); + context->physical_offset = known_host.offset(); + context->physical_capacity = known_host.capacity(); + } + if (known_direct_host) { + hrx_buffer_retain(known_host.buffer()); + context->buffer = known_host.buffer(); + context->direct_host_binding = true; + } + // The importer retains the mapped region, which may be read-only. HRX owns only this descriptor; mutation + // callbacks require caller-writable storage and native weight storage is materialized on demand. + return ggml_backend_buffer_init(&device_context(device)->buft, buffer_i, context, size); +} + static bool eager_capability_declared(enum ggml_op op) { switch (op) { // The scheduler probes preallocated weight tensors as NONE operations when deciding whether their buffer type is @@ -1073,7 +1205,7 @@ static const ggml_backend_device_i device_i = { device_init, device_buffer_type, device_host_buffer_type, - nullptr, + device_buffer_from_host_ptr, device_supports_op, device_supports_buffer_type, nullptr, @@ -1201,13 +1333,33 @@ bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_c if (!ggml_backend_is_hrx(backend) || stats == nullptr) { return false; } - auto * context = static_cast(backend->context); + auto * context = static_cast(backend->context); + bool completed = false; + const ggml::hrx::Status retirement = context->weight_submissions.poll_completed(&completed); + if (!retirement.success()) { + return false; + } + if (completed) { + context->host_weights.evict_except(context->weight_submissions.current_demand(), + context->weight_submissions.outstanding_leases()); + } const ggml::hrx::GraphProgramCacheStats graph_stats = context->graph_programs.stats(); const ggml::hrx::PreparedCommandProgramCacheStats prepared_stats = context->prepared_programs.stats(); + const ggml::hrx::HostWeightCacheStats weight_stats = context->host_weights.stats(); + const ggml::hrx::HostTransferStats transfer_stats = context->host_transfers.stats(); stats->graph_program_builds = graph_stats.builds; stats->graph_program_hits = graph_stats.hits; stats->prepared_program_builds = graph_stats.prepared_program_builds + prepared_stats.builds; stats->prepared_program_hits = graph_stats.prepared_program_hits + prepared_stats.hits; + stats->weight_materializations = weight_stats.misses; + stats->weight_reuses = weight_stats.hits; + stats->weight_uploaded_bytes = weight_stats.materialized_bytes; + stats->weight_live_allocations = weight_stats.allocation_count; + stats->weight_live_bytes = weight_stats.resident_bytes; + stats->weight_peak_bytes = weight_stats.peak_resident_bytes; + stats->host_uploads = transfer_stats.uploads; + stats->host_upload_bytes = transfer_stats.upload_bytes; + stats->weight_duplicate_warnings = weight_stats.duplicate_warnings; return true; } @@ -1222,4 +1374,44 @@ ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device) { nullptr; } +bool ggml_backend_hrx_buffer_is_imported(ggml_backend_buffer_t buffer) { + if (buffer == nullptr) { + return false; + } + const auto * registry = static_cast(ggml_backend_hrx_reg()->context); + for (const auto & device : registry->device_contexts) { + const bool regular_type = buffer->buft == &device->buft; + const bool host_type = buffer->buft == &device->host_buft; + const bool hrx_type = regular_type || host_type; + if (hrx_type) { + return buffer_context(buffer)->imported; + } + } + return false; +} + +void ggml_backend_hrx_mark_imported_region_written(const ggml_tensor * tensor) { + ggml_backend_hrx_buffer_context * context = nullptr; + size_t offset = 0; + if (ggml_backend_hrx_tensor_binding(tensor, &context, &offset)) { + if (context->imported) { + ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); + } + } +} + +void ggml_backend_hrx_mark_retained_region_written(const ggml_tensor * tensor) { + ggml_backend_hrx_buffer_context * context = nullptr; + size_t offset = 0; + if (!ggml_backend_hrx_tensor_binding(tensor, &context, &offset)) { + ggml_backend_hrx_mark_foreign_host_region_written(tensor); + return; + } + ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + const bool retained = context->imported || ggml_backend_buffer_is_host(buffer); + if (retained) { + ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); + } +} + GGML_BACKEND_DL_IMPL(ggml_backend_hrx_reg) diff --git a/ggml/src/ggml-hrx/graph/value-map.cpp b/ggml/src/ggml-hrx/graph/value-map.cpp index dfeea7c027ab..ec39f0eebcfd 100644 --- a/ggml/src/ggml-hrx/graph/value-map.cpp +++ b/ggml/src/ggml-hrx/graph/value-map.cpp @@ -184,6 +184,9 @@ std::optional ValueMap::resolve_buffer_binding(ValueId id) c return std::nullopt; } binding.offset += value->storage_offset; + if (binding.source_host_data != nullptr) { + binding.source_offset += value->storage_offset; + } binding.length = value->byte_count; return binding; } diff --git a/ggml/src/ggml-hrx/graph/value-map.h b/ggml/src/ggml-hrx/graph/value-map.h index b026de65fde6..46af13835a13 100644 --- a/ggml/src/ggml-hrx/graph/value-map.h +++ b/ggml/src/ggml-hrx/graph/value-map.h @@ -63,6 +63,10 @@ struct ValueBufferBinding { size_t capacity = 0; void * host_data = nullptr; bool weight = false; + // Retained host bytes can back a cached native weight even when the binding is directly addressable. + const void * source_host_data = nullptr; + size_t source_offset = 0; + size_t source_capacity = 0; bool requires_materialization() const { return host_data != nullptr; } }; diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 7723246b29f9..06991016e68e 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -59,6 +59,10 @@ void GraphReplayStreamState::add_host_writeback(void * host_destination, pending_host_writebacks_.push_back({ host_destination, mapped_source, size, buffer }); } +bool GraphReplayStreamState::has_pending_host_writebacks() const { + return !pending_host_writebacks_.empty(); +} + PreparedProgramConstantBuffer::~PreparedProgramConstantBuffer() { if (buffer != nullptr) { hrx_buffer_release(buffer); @@ -107,6 +111,7 @@ RecordedCommandGraph::RecordedCommandGraph(RecordedCommandGraph && other) noexce graph(std::exchange(other.graph, nullptr)), exec(std::exchange(other.exec, nullptr)), retained_buffers(std::move(other.retained_buffers)), + resident_host_weights(std::move(other.resident_host_weights)), retained_executables(std::move(other.retained_executables)), bound_transient_arena_allocation_id(other.bound_transient_arena_allocation_id), dispatch_count(other.dispatch_count), @@ -132,6 +137,7 @@ RecordedCommandGraph & RecordedCommandGraph::operator=(RecordedCommandGraph && o graph = std::exchange(other.graph, nullptr); exec = std::exchange(other.exec, nullptr); retained_buffers = std::move(other.retained_buffers); + resident_host_weights = std::move(other.resident_host_weights); retained_executables = std::move(other.retained_executables); bound_transient_arena_allocation_id = other.bound_transient_arena_allocation_id; dispatch_count = other.dispatch_count; @@ -699,18 +705,9 @@ static Dispatch build_dispatch(const ResolvedCommand & command) { } struct GraphValueAccess { - bool read = false; - bool write = false; + bool read = false; + bool write = false; bool main_write = false; - bool native_read = false; - bool transformed_read = false; - bool layout_conflict = false; - std::string layout = kNativeWeightLayout; - ggml_type source_type = GGML_TYPE_COUNT; - int64_t input_size = 0; - int64_t output_size = 0; - size_t source_length = 0; - size_t materialized_length = 0; }; static std::unordered_map collect_graph_value_access(const CommandProgram & commands) { @@ -735,38 +732,6 @@ static std::unordered_map collect_graph_value_access( access.write = true; break; } - if (binding.access == ResourceAccess::Write) { - continue; - } - if (binding.layout == kNativeWeightLayout) { - access.native_read = true; - if (access.transformed_read) { - access.layout_conflict = true; - } - continue; - } - if (binding.offset != 0 || binding.source_length == 0 || binding.length == 0 || - binding.source_type == GGML_TYPE_COUNT || binding.input_size <= 0 || binding.output_size <= 0) { - access.layout_conflict = true; - continue; - } - if (!access.transformed_read) { - access.transformed_read = true; - access.layout = binding.layout; - access.source_type = binding.source_type; - access.input_size = binding.input_size; - access.output_size = binding.output_size; - access.source_length = binding.source_length; - access.materialized_length = binding.length; - } else if (access.layout != binding.layout || access.source_type != binding.source_type || - access.input_size != binding.input_size || access.output_size != binding.output_size || - access.source_length != binding.source_length || - access.materialized_length != binding.length) { - access.layout_conflict = true; - } - if (access.native_read) { - access.layout_conflict = true; - } } } }; @@ -775,6 +740,50 @@ static std::unordered_map collect_graph_value_access( return access_by_value; } +static bool same_weight_request(const CommandBinding & lhs, const CommandBinding & rhs) { + const bool same_value = lhs.value == rhs.value; + const bool same_layout = lhs.layout == rhs.layout; + const bool native_layout = lhs.layout == kNativeWeightLayout; + const bool same_transformation = native_layout || + (lhs.source_type == rhs.source_type && lhs.input_size == rhs.input_size && + lhs.output_size == rhs.output_size && lhs.source_length == rhs.source_length && lhs.length == rhs.length); + return same_value && same_layout && same_transformation; +} + +static void append_weight_requests(const std::vector & commands, + const std::unordered_map & access_by_value, + std::unordered_map> & groups_by_value, + CommandProgramWeightRequests & result) { + for (const Command & command : commands) { + for (const CommandBinding & request : command.bindings) { + const bool graph_value = request.origin == CommandBindingOrigin::GraphValue; + const bool reads = request.access != ResourceAccess::Write; + const bool eligible_access = graph_value && reads; + if (!eligible_access) { + continue; + } + const auto found_access = access_by_value.find(request.value.value); + const bool value_has_access = found_access != access_by_value.end(); + const bool value_is_written = value_has_access && found_access->second.write; + if (value_is_written) { + continue; + } + auto & groups = groups_by_value[request.value.value]; + const auto found_group = std::find_if(groups.begin(), groups.end(), [&](size_t index) { + return same_weight_request(result.requests[index].binding, request); + }); + size_t group; + if (found_group == groups.end()) { + group = result.group_count++; + groups.push_back(result.requests.size()); + } else { + group = result.requests[*found_group].group; + } + result.requests.push_back({ request, group }); + } + } +} + static CommandProgramBindings materialize_host_bindings(const CommandProgramExecutionContext & context, const CommandProgram & commands, const CommandProgramBindings & bindings, @@ -783,57 +792,52 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec Status status; materialized.reserve(bindings.bindings().size()); const std::unordered_map access_by_value = collect_graph_value_access(commands); + const CommandProgramWeightDemands demands = collect_command_program_weight_demands(commands, bindings); + status.append(demands.status); for (const CommandProgramBinding & binding : bindings.bindings()) { const auto found_access = access_by_value.find(binding.value.value); const GraphValueAccess access = found_access != access_by_value.end() ? found_access->second : GraphValueAccess{}; - const bool materialize_weight = binding.weight && access.read && !access.write && - (binding.requires_materialization() || access.transformed_read); - if (binding.length == 0 || (!binding.requires_materialization() && !materialize_weight)) { + const bool managed_weight = binding.weight && access.read && !access.write && + binding.source_host_data != nullptr; + const bool empty_binding = binding.length == 0; + const bool directly_bindable = !binding.requires_materialization() && !managed_weight; + // Elided operations can leave external weights with no command users. + const bool unused_binding = !access.read && !access.write; + const bool use_existing_binding = empty_binding || directly_bindable || unused_binding; + if (use_existing_binding) { materialized.push_back(binding); continue; } - if (materialize_weight) { - if (access.layout_conflict) { - status.log("weight value %d has conflicting resident layout requests", binding.value.value); - materialized.push_back(binding); - continue; - } - HostWeightSource source; - source.host_data = binding.host_data; - source.device_buffer = binding.host_data == nullptr ? binding.buffer : nullptr; - source.identity = binding.identity; - source.generation = binding.generation; - source.capacity = binding.capacity; - source.offset = binding.offset; - source.length = access.transformed_read ? access.source_length : binding.length; - source.materialized_length = access.transformed_read ? access.materialized_length : binding.length; - source.layout = access.layout; - source.source_type = access.source_type; - source.input_size = access.input_size; - source.output_size = access.output_size; - if (source.length > binding.length) { - status.log("weight value %d layout source length %zu exceeds runtime length %zu", binding.value.value, - source.length, binding.length); - materialized.push_back(binding); - continue; - } - HostWeightAcquireResult resident = - context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source); - if (!resident.valid()) { - status.log("materialize host weight value %d failed", binding.value.value); - status.append(resident.status); - materialized.push_back(binding); - continue; + if (managed_weight) { + for (const CommandProgramWeightDemand & demand : demands.weights) { + if (demand.request.value != binding.value) { + continue; + } + CommandProgramBinding device_binding = binding; + device_binding.layout = demand.request.layout; + device_binding.layout_source_type = demand.request.source_type; + device_binding.input_size = demand.request.input_size; + device_binding.output_size = demand.request.output_size; + device_binding.source_length = demand.request.source_length; + if (demand.requires_materialization) { + HostWeightAcquireResult resident = context.host_weights->acquire( + context.device, context.stream, *context.host_transfers, demand.source); + if (!resident.valid()) { + status.log("materialize host weight value %d layout %s failed", binding.value.value, + demand.request.layout.c_str()); + status.append(resident.status); + continue; + } + device_binding.buffer = resident.lease.buffer(); + device_binding.host_data = nullptr; + device_binding.offset = 0; + device_binding.length = resident.lease.length(); + device_binding.capacity = resident.lease.length(); + prepared.resident_host_weights.push_back(std::move(resident.lease)); + } + materialized.push_back(std::move(device_binding)); } - CommandProgramBinding device_binding = binding; - device_binding.buffer = resident.lease.buffer(); - device_binding.host_data = nullptr; - device_binding.offset = 0; - device_binding.length = resident.lease.length(); - device_binding.capacity = resident.lease.length(); - materialized.push_back(device_binding); - prepared.resident_host_weights.push_back(std::move(resident.lease)); continue; } @@ -866,6 +870,11 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec materialized.push_back(device_binding); prepared.host_staging.push_back(std::move(staging)); } + const bool can_report_residency = status.success() && context.host_weights != nullptr; + if (can_report_residency) { + context.host_weights->report_coexisting_representations(prepared.resident_host_weights, + "same command program"); + } return CommandProgramBindings::from_bindings(std::move(materialized), status); } @@ -1647,6 +1656,8 @@ static void retain_prepared_command_list_resources(const std::vector retained_buffers; std::unordered_map retained_executables; if (allocation.buffer != nullptr) { @@ -1778,6 +1789,91 @@ static bool execute_prepared_command_prefix_via_graph(const CommandProgramExecut } // namespace +CommandProgramWeightRequests build_command_program_weight_requests(const CommandProgram & commands) { + CommandProgramWeightRequests result; + std::unordered_map> groups_by_value; + const auto access_by_value = collect_graph_value_access(commands); + append_weight_requests(commands.initialization_commands, access_by_value, groups_by_value, result); + append_weight_requests(commands.commands, access_by_value, groups_by_value, result); + return result; +} + +CommandProgramWeightDemands bind_command_program_weight_demands(const CommandProgramWeightRequests & requests, + const CommandProgramBindings & bindings) { + CommandProgramWeightDemands result; + std::unordered_map bindings_by_value; + bindings_by_value.reserve(bindings.bindings().size()); + for (const CommandProgramBinding & binding : bindings.bindings()) { + bindings_by_value.try_emplace(binding.value.value, &binding); + } + std::vector emitted_groups(requests.group_count, 0); + result.weights.reserve(requests.group_count); + for (const CommandProgramWeightRequest & entry : requests.requests) { + const CommandBinding & request = entry.binding; + const auto found_binding = bindings_by_value.find(request.value.value); + const CommandProgramBinding * binding = + found_binding != bindings_by_value.end() ? found_binding->second : nullptr; + const bool bound_weight = binding != nullptr && binding->weight; + if (!bound_weight) { + continue; + } + const bool transformed = request.layout != kNativeWeightLayout; + const bool retained_source = binding->source_host_data != nullptr; + const bool missing_transformed_source = transformed && !retained_source; + if (missing_transformed_source) { + result.status.log("weight value %d has no retained source for transformed layout %s", + request.value.value, request.layout.c_str()); + continue; + } + if (!retained_source) { + continue; + } + if (transformed) { + const bool valid_shape = request.source_type != GGML_TYPE_COUNT && request.input_size > 0 && + request.output_size > 0; + const bool valid_range = request.offset == 0 && request.source_length != 0 && + request.source_length <= binding->length && request.length != 0; + const bool valid_transformation = valid_shape && valid_range; + if (!valid_transformation) { + result.status.log("weight value %d has invalid transformed layout request %s", + request.value.value, request.layout.c_str()); + continue; + } + } + // Validation precedes deduplication: an invalid duplicate must still report its error, + // and an invalid first occurrence must not suppress a later valid request. + if (emitted_groups[entry.group]) { + continue; + } + emitted_groups[entry.group] = 1; + CommandProgramWeightDemand demand; + demand.request = request; + demand.requires_materialization = binding->requires_materialization() || transformed; + demand.source.host_data = binding->source_host_data; + demand.source.identity = binding->identity; + demand.source.generation = binding->generation; + demand.source.capacity = binding->source_capacity; + demand.source.offset = binding->source_offset; + demand.source.length = transformed ? request.source_length : binding->length; + demand.source.region_length = binding->length; + demand.source.region_type = binding->source_type; + demand.source.materialized_length = transformed ? request.length : binding->length; + demand.source.layout = request.layout; + demand.source.source_type = transformed ? request.source_type : binding->source_type; + demand.source.input_size = request.input_size; + demand.source.output_size = request.output_size; + demand.source.source_ne = binding->source_ne; + demand.source.source_nb = binding->source_nb; + result.weights.push_back(std::move(demand)); + } + return result; +} + +CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, + const CommandProgramBindings & bindings) { + return bind_command_program_weight_demands(build_command_program_weight_requests(commands), bindings); +} + bool debug_serial_command_execution_enabled() { return environment_flag_enabled("GGML_HRX_DEBUG_SERIAL_EXECUTION") || dispatch_trace_enabled(); } diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index bc1debfe03a3..211953fa5b4a 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -45,6 +45,7 @@ struct GraphReplayStreamState { void mark_stream_synchronized(); void add_host_writeback(void * host_destination, const void * mapped_source, size_t size, hrx_buffer_t buffer); + bool has_pending_host_writebacks() const; void clear(); private: @@ -64,6 +65,38 @@ struct CommandProgramExecutionContext { GraphReplayStreamState * graph_replay_state = nullptr; }; +// Requests contain command metadata only. Every occurrence remains available for validation before +// its deduplication group is emitted, since same-layout requests can differ in validity. +struct CommandProgramWeightRequest { + CommandBinding binding; + size_t group = 0; +}; + +struct CommandProgramWeightRequests { + std::vector requests; + size_t group_count = 0; +}; + +struct CommandProgramWeightDemand { + CommandBinding request; + HostWeightSource source; + bool requires_materialization = false; +}; + +struct CommandProgramWeightDemands { + std::vector weights; + Status status; + + bool valid() const { return status.success(); } +}; + +CommandProgramWeightRequests build_command_program_weight_requests(const CommandProgram & commands); +CommandProgramWeightDemands bind_command_program_weight_demands(const CommandProgramWeightRequests & requests, + const CommandProgramBindings & bindings); + +CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, + const CommandProgramBindings & bindings); + struct PreparedCommandBinding { CommandBinding binding; ResolvedBufferRef ref; @@ -115,6 +148,7 @@ struct RecordedCommandGraph { hrx_graph_t graph = nullptr; hrx_graph_exec_t exec = nullptr; std::vector retained_buffers; + std::vector resident_host_weights; std::vector retained_executables; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; size_t dispatch_count = 0; @@ -132,6 +166,13 @@ struct RecordedCommandGraph { bool valid() const { return status.success() && exec != nullptr; } }; +// A submitted command may still use prepared bindings after its cache entry is retired. +// The execution completion tracker retains this object until the stream reaches the submission. +struct CachedCommandProgramResources { + PreparedCommandProgram prepared; + RecordedCommandGraph recorded; +}; + struct RecordedCommandGraphExecutionResult { bool success = false; Status status; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index fa509c839792..4bbd4b3eaadf 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -2,15 +2,285 @@ #include "backend-buffer-binding.h" #include "ggml-impl.h" +#include "hrx-interop-utils.h" +#include "hrx_runtime.h" #include "runtime/kernel-executable-cache.h" #include "runtime/prepared-command-program-cache.h" #include "runtime/transient-arena.h" +#include +#include #include #include namespace ggml::hrx { +WeightSubmissionTracker::~WeightSubmissionTracker() { + mark_stream_synchronized(); +} + +bool WeightSubmissionTracker::RegionKey::operator==(const RegionKey & other) const { + return identity == other.identity && offset == other.offset && + length == other.length && type == other.type && ne == other.ne && nb == other.nb; +} + +size_t WeightSubmissionTracker::RegionKeyHash::operator()(const RegionKey & key) const { + uint64_t hash = UINT64_C(1469598103934665603); + auto mix = [&](uint64_t value) { + hash ^= value; + hash *= UINT64_C(1099511628211); + }; + mix(key.identity); + mix(key.offset); + mix(key.length); + mix(static_cast(key.type)); + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + mix(static_cast(key.ne[dimension])); + mix(key.nb[dimension]); + } + return static_cast(hash); +} + +WeightSubmissionTracker::RegionKey WeightSubmissionTracker::region_of(const HostWeightSource & source) { + return { source.identity, source.offset, source.region_length, source.region_type, + source.source_ne, source.source_nb }; +} + +static bool same_weight_representation(const HostWeightSource & left, const HostWeightSource & right) { + return left.generation == right.generation && left.capacity == right.capacity && + left.length == right.length && + left.materialized_length == right.materialized_length && left.layout == right.layout && + left.source_type == right.source_type && left.input_size == right.input_size && + left.output_size == right.output_size; +} + +WeightSubmissionTracker::RegionDemands WeightSubmissionTracker::group_demands( + const std::vector & demand) { + RegionDemands grouped; + for (const HostWeightSource & source : demand) { + auto & layouts = grouped[region_of(source)]; + const bool duplicate = std::any_of(layouts.begin(), layouts.end(), [&](const HostWeightSource & existing) { + return same_weight_representation(existing, source); + }); + if (!duplicate) { + layouts.push_back(source); + } + } + return grouped; +} + +bool WeightSubmissionTracker::demand_matches(const std::vector & demand) const { + bool needs_grouped_comparison = false; + for (const HostWeightSource & source : demand) { + const auto retained = retained_regions_.find(region_of(source)); + if (retained == retained_regions_.end()) { + return false; + } + const bool singleton_region = retained->second.size() == 1; + if (!singleton_region) { + needs_grouped_comparison = true; + break; + } + if (!same_weight_representation(retained->second.front(), source)) { + return false; + } + } + if (!needs_grouped_comparison) { + return true; + } + + const RegionDemands touched = group_demands(demand); + for (const auto & [region, layouts] : touched) { + const auto retained = retained_regions_.find(region); + const bool region_is_retained = retained != retained_regions_.end(); + if (!region_is_retained) { + return false; + } + const bool same_layout_count = layouts.size() == retained->second.size(); + if (!same_layout_count) { + return false; + } + for (const HostWeightSource & source : layouts) { + const bool present = std::any_of(retained->second.begin(), retained->second.end(), + [&](const HostWeightSource & existing) { + return same_weight_representation(existing, source); + }); + if (!present) { + return false; + } + } + } + return true; +} + +std::vector WeightSubmissionTracker::merged_demand_for( + const std::vector & demand) const { + RegionDemands merged = retained_regions_; + for (auto & [region, layouts] : group_demands(demand)) { + merged[region] = std::move(layouts); + } + std::vector result; + for (const auto & [region, layouts] : merged) { + result.insert(result.end(), layouts.begin(), layouts.end()); + } + return result; +} + +void WeightSubmissionTracker::publish_demand(const std::vector & demand) { + for (auto & [region, layouts] : group_demands(demand)) { + retained_regions_[region] = std::move(layouts); + } + current_demand_.clear(); + for (const auto & [region, layouts] : retained_regions_) { + current_demand_.insert(current_demand_.end(), layouts.begin(), layouts.end()); + } +} + +Status WeightSubmissionTracker::poll_completed(bool * completed) { + Status status; + if (completed != nullptr) { + *completed = false; + } + for (auto pending = pending_.begin(); pending != pending_.end();) { + const bool has_event = pending->event != nullptr; + const bool has_timeline = pending->semaphore != nullptr; + const bool has_completion = has_event || has_timeline; + if (!has_completion) { + ++pending; + continue; + } + bool complete = false; + if (has_event) { + if (ErrorResult error = take_status(hrx_event_query(pending->event, &complete))) { + status.log("query HRX weight submission completion: %s", error->c_str()); + return status; + } + } else { + uint64_t completed_value = 0; + if (ErrorResult error = take_status(hrx_semaphore_query(pending->semaphore, &completed_value))) { + status.log("query HRX weight submission timeline: %s", error->c_str()); + return status; + } + complete = completed_value >= pending->value; + } + if (complete) { + const bool report_obsolete_completion = completed != nullptr && pending->may_hold_obsolete_demand; + if (report_obsolete_completion) { + *completed = true; + } + release_completion(*pending); + pending = pending_.erase(pending); + } else { + ++pending; + } + } + return status; +} + +void WeightSubmissionTracker::release_completion(Pending & pending) { + if (pending.event != nullptr) { + hrx_event_release(pending.event); + pending.event = nullptr; + } + if (pending.semaphore != nullptr) { + hrx_semaphore_release(pending.semaphore); + pending.semaphore = nullptr; + } +} + +Status WeightSubmissionTracker::record(hrx_device_t device, hrx_stream_t stream, + std::shared_ptr resources, + const std::vector & demand, bool admitted, + bool demand_changed) { + Status status; + if (resources == nullptr) { + if (admitted) { + if (demand_changed) { + publish_demand(demand); + } + current_leases_.clear(); + } + return status; + } + const bool admitted_new_demand = admitted && demand_changed; + if (admitted_new_demand) { + for (Pending & pending : pending_) { + pending.may_hold_obsolete_demand = true; + } + } + hrx_event_t event = nullptr; + hrx_semaphore_t semaphore = nullptr; + uint64_t value = 0; + if (fail_next_record_) { + fail_next_record_ = false; + status.log("injected HRX weight completion capture failure"); + } else if (admitted) { + // Capture after all queued work, including a possible host-download tail. + hrx_timeline_point_t position = {}; + if (ErrorResult error = take_status(hrx_stream_flush(stream))) { + status.log("flush HRX weight submission: %s", error->c_str()); + } else if (ErrorResult error = take_status(hrx_stream_get_timeline_position(stream, &position))) { + status.log("capture HRX weight submission timeline: %s", error->c_str()); + } else { + semaphore = position.semaphore; + value = position.value; + hrx_semaphore_retain(semaphore); + } + } else if (ErrorResult error = take_status(hrx_event_create(device, HRX_EVENT_FLAG_DISABLE_TIMING, &event))) { + status.log("create HRX weight submission completion: %s", error->c_str()); + } else if (ErrorResult error = take_status(hrx_event_record(event, stream))) { + status.log("record HRX weight submission completion: %s", error->c_str()); + } + if (!status.success()) { + if (event != nullptr) { + hrx_event_release(event); + } + if (semaphore != nullptr) { + hrx_semaphore_release(semaphore); + } + // A capture failure must not let queued commands outlive their bindings. + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize after HRX completion capture failure: %s", error->c_str()); + pending_.push_back({ nullptr, nullptr, 0, std::move(resources), true }); + return status; + } + mark_stream_synchronized(); + return status; + } + pending_.push_back({ event, semaphore, value, std::move(resources), !admitted }); + if (admitted) { + if (admitted_new_demand) { + publish_demand(demand); + } + current_leases_ = pending_.back().resources->prepared.resident_host_weights; + } + return status; +} + +void WeightSubmissionTracker::mark_stream_synchronized() { + for (Pending & pending : pending_) { + release_completion(pending); + } + pending_.clear(); +} + +bool WeightSubmissionTracker::has_obsolete_pending() const { + return std::any_of(pending_.begin(), pending_.end(), + [](const Pending & pending) { return pending.may_hold_obsolete_demand; }); +} + +std::vector WeightSubmissionTracker::outstanding_leases() const { + std::vector leases; + for (const Pending & pending : pending_) { + if (pending.resources == nullptr) { + continue; + } + const auto & resident = pending.resources->prepared.resident_host_weights; + leases.insert(leases.end(), resident.begin(), resident.end()); + } + return leases; +} + GraphExecutor::GraphExecutor(ggml_backend_hrx_context & context) : context_(context) {} Status GraphExecutor::context_valid_for_graph_programs() const { @@ -54,15 +324,14 @@ CommandProgramBindings GraphExecutor::bind_external_value_buffers(const GraphPro CommandProgramBinding binding; binding.value = external.value; if (ggml_backend_hrx_resolve_value_buffer(external.tensor, value_binding)) { - binding.buffer = value_binding.buffer; - binding.host_data = value_binding.host_data; - binding.offset = value_binding.offset; - binding.length = value_binding.length; - binding.identity = value_binding.identity; - binding.generation = value_binding.generation; - binding.capacity = value_binding.capacity; - binding.weight = value_binding.weight; - binding.empty_value = ggml_nbytes(external.tensor) == 0; + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = external.tensor->ne[dimension]; + nb[dimension] = external.tensor->nb[dimension]; + } + binding = make_command_program_binding(external.value, value_binding, external.tensor->type, ne, nb, + ggml_nbytes(external.tensor) == 0); } else { status.log("external value %d is not bound", external.value.value); } @@ -95,12 +364,18 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); GraphProgramMatch binding_match = std::move(lookup.match); - if (use_graph_prepared && lookup.program->has_prepared_program()) { - binding_match = lookup.program->match_host_staging_graph(graph); - if (!binding_match.valid()) { - result.status.append(binding_match.status); - return result; - } + const bool has_prepared = lookup.program->has_prepared_program(); + const bool may_use_partial_bindings = use_graph_prepared && has_prepared; + bool full_bindings = true; + if (may_use_partial_bindings) { + const bool imported_sources_current = lookup.program->imported_sources_current(); + binding_match = imported_sources_current ? lookup.program->match_host_staging_graph(graph) : + lookup.program->match_trusted_graph(graph, true); + full_bindings = !imported_sources_current; + } + if (!binding_match.valid()) { + result.status.append(binding_match.status); + return result; } CommandProgramBindings bindings = bind_external_value_buffers(binding_match); @@ -108,6 +383,90 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { result.status.append(bindings.status); return result; } + std::shared_ptr> cached_demand; + if (!full_bindings) { + const GraphProgramMatch candidate_match = lookup.program->match_weight_demand_graph(graph); + if (!candidate_match.valid()) { + result.status.append(candidate_match.status); + return result; + } + CommandProgramBindings candidate_bindings = bind_external_value_buffers(candidate_match); + if (!candidate_bindings.valid()) { + result.status.append(candidate_bindings.status); + return result; + } + cached_demand = lookup.program->cached_weight_demand_for(candidate_bindings); + } + std::optional full_demand_bindings; + const CommandProgramBindings * demand_bindings = &bindings; + const bool needs_full_demand_bindings = !full_bindings && cached_demand == nullptr; + if (needs_full_demand_bindings) { + const GraphProgramMatch demand_match = lookup.program->match_trusted_graph(graph, true); + if (!demand_match.valid()) { + result.status.append(demand_match.status); + return result; + } + full_demand_bindings.emplace(bind_external_value_buffers(demand_match)); + if (!full_demand_bindings->valid()) { + result.status.append(full_demand_bindings->status); + return result; + } + demand_bindings = &*full_demand_bindings; + } + std::vector collected_demand; + if (cached_demand == nullptr) { + const CommandProgramWeightDemands weight_demands = + bind_command_program_weight_demands(lookup.program->weight_requests(), *demand_bindings); + if (!weight_demands.valid()) { + result.status.append(weight_demands.status); + return result; + } + for (const CommandProgramWeightDemand & demand : weight_demands.weights) { + collected_demand.push_back(demand.source); + } + } + const std::vector & current_demand = + cached_demand != nullptr ? *cached_demand : collected_demand; + const bool demand_changed = !context_.weight_submissions.demand_matches(current_demand); + std::vector merged_demand; + if (demand_changed) { + merged_demand = context_.weight_submissions.merged_demand_for(current_demand); + } + bool completed_before = false; + result.status = context_.weight_submissions.poll_completed(&completed_before); + if (!result.status.success()) { + return result; + } + if (demand_changed) { + context_.weight_submissions.release_current_leases(); + context_.graph_programs.invalidate_obsolete_weights(merged_demand); + context_.prepared_programs.invalidate_obsolete_weights(merged_demand); + } + const bool prepared_program_missing = !lookup.program->has_prepared_program(); + const bool restore_full_bindings = use_graph_prepared && prepared_program_missing; + if (restore_full_bindings) { + if (full_demand_bindings) { + bindings = std::move(*full_demand_bindings); + full_bindings = true; + } else if (!full_bindings) { + const GraphProgramMatch full_match = lookup.program->match_trusted_graph(graph, true); + if (!full_match.valid()) { + result.status.append(full_match.status); + return result; + } + bindings = bind_external_value_buffers(full_match); + if (!bindings.valid()) { + result.status.append(bindings.status); + return result; + } + full_bindings = true; + } + } + const bool reconcile_before = demand_changed || completed_before; + if (reconcile_before) { + const auto & retained = demand_changed ? merged_demand : context_.weight_submissions.current_demand(); + context_.host_weights.evict_except(retained, context_.weight_submissions.outstanding_leases()); + } const CommandProgramExecutionContext execution_context = { context_.device->device, context_.stream, @@ -121,10 +480,56 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { &context_.graph_replay_state, }; const PreparedCommandProgramCacheExecutionResult execution = - use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings) : + use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings, full_bindings) : context_.prepared_programs.execute_with_result(execution_context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), bindings); + result.graph_replay_event = execution.graph_replay_event; + const Status completion_status = context_.weight_submissions.record( + context_.device->device, context_.stream, execution.execution_resources, current_demand, execution.success, + demand_changed); + result.status.append(completion_status); + const bool admitted = execution.success && completion_status.success(); + const bool cache_needs_refresh = cached_demand == nullptr || restore_full_bindings; + const bool cache_can_refresh = admitted && use_graph_prepared && cache_needs_refresh; + if (cache_can_refresh) { + const CommandProgramBindings & cache_bindings = + full_bindings ? bindings : *full_demand_bindings; + lookup.program->remember_weight_demand(cache_bindings, current_demand); + } + if (!admitted) { + const auto & retained_demand = context_.weight_submissions.current_demand(); + context_.graph_programs.invalidate_obsolete_weights(retained_demand); + context_.prepared_programs.invalidate_obsolete_weights(retained_demand); + } + bool completed_after = false; + const Status retirement_status = context_.weight_submissions.poll_completed(&completed_after); + result.status.append(retirement_status); + const bool reconcile_after = demand_changed || completed_after || !admitted; + if (reconcile_after) { + const std::vector outstanding = context_.weight_submissions.outstanding_leases(); + context_.host_weights.evict_except(context_.weight_submissions.current_demand(), outstanding); + const auto & current = context_.weight_submissions.current_leases(); + const bool outstanding_only_allocation = std::any_of(outstanding.begin(), outstanding.end(), + [&](const HostWeightLease & lease) { + return std::none_of(current.begin(), current.end(), + [&](const HostWeightLease & active) { return active.buffer() == lease.buffer(); }); + }); + const bool report_overlap = admitted && outstanding_only_allocation; + if (report_overlap) { + std::vector coexisting = current; + coexisting.insert(coexisting.end(), outstanding.begin(), outstanding.end()); + context_.host_weights.report_coexisting_representations(coexisting, + "current_vs_outstanding"); + } + } + if (execution.execution_resources != nullptr) { + // A failed return can still follow queued writes or a successful fallback synchronization. + lookup.program->mark_imported_writes(graph); + } + if (!result.status.success()) { + return result; + } if (!execution.success) { result.status.append(execution.status); if (result.status.success()) { diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.h b/ggml/src/ggml-hrx/runtime/graph-executor.h index a2759ea71875..7b66ae160ccc 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.h +++ b/ggml/src/ggml-hrx/runtime/graph-executor.h @@ -20,6 +20,7 @@ struct GraphSupportResult { struct GraphExecutionResult { enum ggml_status code = GGML_STATUS_FAILED; Status status; + HrxGraphReplayEvent graph_replay_event = HrxGraphReplayEvent::Disabled; bool success() const { return code == GGML_STATUS_SUCCESS && status.success(); } }; diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index 666602f8fcd3..f63d4255561f 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -1,13 +1,18 @@ #include "graph-program-cache.h" +#include "backend-buffer-binding.h" #include "dispatch/command-program-dump.h" #include "dispatch/dispatch-scheduler.h" #include "ggml-impl.h" #include "ggml.h" +#include "hrx-interop-utils.h" +#include "hrx_runtime.h" +#include #include #include #include +#include #include #include #include @@ -15,6 +20,31 @@ namespace ggml::hrx { namespace { +static bool same_weight_demand_binding(const CommandProgramBinding & left, const CommandProgramBinding & right) { + const bool same_value = left.value == right.value; + const bool same_weight_class = left.weight == right.weight; + const bool same_binding_class = same_value && same_weight_class; + if (!same_binding_class) { + return false; + } + if (!left.weight) { + // Mutable graph inputs can change each token without changing which weights are demanded. + return true; + } + const bool same_storage = left.buffer == right.buffer && left.host_data == right.host_data && + left.offset == right.offset && left.length == right.length && + left.capacity == right.capacity && left.empty_value == right.empty_value; + const bool same_source = left.identity == right.identity && left.generation == right.generation && + left.source_host_data == right.source_host_data && + left.source_offset == right.source_offset && left.source_capacity == right.source_capacity && + left.source_type == right.source_type && left.source_ne == right.source_ne && + left.source_nb == right.source_nb; + const bool same_layout = left.layout == right.layout && left.layout_source_type == right.layout_source_type && + left.input_size == right.input_size && left.output_size == right.output_size && + left.source_length == right.source_length; + return same_storage && same_source && same_layout; +} + static const ggml_tensor * tensor_storage_root(const ggml_tensor * tensor) { while (tensor != nullptr && tensor->view_src != nullptr) { tensor = tensor->view_src; @@ -322,6 +352,7 @@ GraphProgram::GraphProgram(uint64_t uid, target_(std::move(target)), graph_(std::move(graph)), commands_(std::move(commands)), + weight_requests_(build_command_program_weight_requests(*commands_)), command_shape_(std::move(command_shape)), command_shape_hash_(command_program_shape_hash(command_shape_)) {} @@ -399,6 +430,16 @@ GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_ } } if (!bind_external) { + for (const GraphProgramExternalSlot & slot : external_slots_) { + const ggml_tensor * tensor = resolve_external_slot(current_graph, slot, result.status); + if (tensor == nullptr) { + return result; + } + if (ggml_backend_hrx_weight_layout_eligible(tensor) != slot.weight_layout_eligible) { + result.status.log("external weight source class changed for value %d", slot.value.value); + return result; + } + } return result; } result.external_bindings.reserve(external_slots_.size()); @@ -407,6 +448,10 @@ GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_ if (tensor == nullptr) { return result; } + if (ggml_backend_hrx_weight_layout_eligible(tensor) != slot.weight_layout_eligible) { + result.status.log("external weight source class changed for value %d", slot.value.value); + return result; + } result.external_bindings.push_back({ slot.value, tensor }); } return result; @@ -419,8 +464,8 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur result.status.log("missing prepared HRX command program"); return result; } - result.external_bindings.reserve(prepared_.host_staging.size()); - for (const HostStagingBuffer & staging : prepared_.host_staging) { + result.external_bindings.reserve(resources_->prepared.host_staging.size()); + for (const HostStagingBuffer & staging : resources_->prepared.host_staging) { const GraphProgramExternalSlot * slot = find_external_slot(ValueId(staging.value)); if (slot == nullptr) { result.status.log("prepared host staging value %d has no external graph slot", staging.value); @@ -435,15 +480,75 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur return result; } +GraphProgramMatch GraphProgram::match_weight_demand_graph(const ggml_cgraph & current_graph) const { + GraphProgramMatch result; + result.external_bindings.reserve(weight_demand_slots_.size()); + for (const GraphProgramExternalSlot & slot : weight_demand_slots_) { + const ggml_tensor * tensor = resolve_external_slot(current_graph, slot, result.status); + if (tensor == nullptr) { + return result; + } + const bool source_class_current = + ggml_backend_hrx_weight_layout_eligible(tensor) == slot.weight_layout_eligible; + if (!source_class_current) { + result.status.log("external weight source class changed for value %d", slot.value.value); + return result; + } + result.external_bindings.push_back({ slot.value, tensor }); + } + return result; +} + Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { Status status; external_slots_.clear(); + weight_demand_slots_.clear(); + weight_demand_values_.clear(); + weight_demand_bindings_.clear(); + weight_demand_.reset(); + weight_demand_resources_.reset(); + imported_source_watches_.clear(); + written_values_.clear(); external_slot_by_value_.clear(); fast_path_nodes_ = graph.nodes; external_slots_.reserve(match.external_bindings.size()); + std::unordered_set read_values; + auto collect_access = [&](const std::vector & commands) { + for (const Command & command : commands) { + for (const CommandBinding & binding : command.bindings) { + const bool graph_value = binding.origin == CommandBindingOrigin::GraphValue; + const bool writes_only = binding.access == ResourceAccess::Write; + const bool reads_and_writes = binding.access == ResourceAccess::ReadWrite; + const bool writes_value = writes_only || reads_and_writes; + const bool graph_write = graph_value && writes_value; + if (graph_write) { + written_values_.insert(binding.value.value); + } + const bool reads_value = binding.access != ResourceAccess::Write; + const bool graph_read = graph_value && reads_value; + if (graph_read) { + read_values.insert(binding.value.value); + } + } + } + }; + collect_access(commands_->initialization_commands); + collect_access(commands_->commands); for (const GraphProgramExternalBinding & binding : match.external_bindings) { + ValueBufferBinding value_binding; + const bool bound = ggml_backend_hrx_resolve_value_buffer(binding.tensor, value_binding); + const bool weight = bound && value_binding.weight; + const bool host_source = bound && value_binding.source_host_data != nullptr; + const bool tracks_generation = weight && host_source; + ggml_backend_hrx_imported_region region = {}; + if (ggml_backend_hrx_get_retained_region(binding.tensor, region)) { + // Watch each physical region, including aliases in one buffer, but only weight contents. + imported_source_watches_.push_back({ binding.tensor, region.identity, region.generation, + region.offset, region.length, tracks_generation }); + } GraphProgramExternalSlot slot; slot.value = binding.value; + slot.weight_layout_eligible = ggml_backend_hrx_weight_layout_eligible(binding.tensor); bool found = false; for (int i = 0; i < graph.n_nodes && !found; ++i) { const ggml_tensor * node = graph.nodes[i]; @@ -476,6 +581,15 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra external_slot_by_value_[binding.value.value] = external_slots_.size(); external_slots_.push_back(slot); } + for (const GraphProgramExternalSlot & slot : external_slots_) { + const bool read_by_command = read_values.count(slot.value.value) != 0; + const bool written_by_command = written_values_.count(slot.value.value) != 0; + const bool candidate = read_by_command && !written_by_command; + if (candidate) { + weight_demand_slots_.push_back(slot); + weight_demand_values_.insert(slot.value.value); + } + } return status; } @@ -489,14 +603,134 @@ bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { return graph.uid == uid_ && graph.nodes == fast_path_nodes_; } +bool GraphProgram::imported_sources_current() const { + std::lock_guard lock(prepared_mutex_); + for (const ImportedSourceWatch & watch : imported_source_watches_) { + ggml_backend_hrx_imported_region region = {}; + if (!ggml_backend_hrx_get_retained_region(watch.tensor, region)) { + return false; + } + const bool same_identity = region.identity == watch.identity; + const bool same_offset = region.offset == watch.offset; + const bool same_length = region.length == watch.length; + const bool same_generation = !watch.tracks_content_generation || region.generation == watch.generation; + const bool identity_changed = !same_identity; + const bool offset_changed = !same_offset; + const bool length_changed = !same_length; + const bool generation_changed = !same_generation; + const bool source_changed = identity_changed || offset_changed || length_changed || generation_changed; + if (source_changed) { + return false; + } + } + return true; +} + +void GraphProgram::mark_imported_writes(const ggml_cgraph & graph) const { + std::unordered_set marked_identities; + for (const GraphProgramExternalSlot & slot : external_slots_) { + const bool written_value = written_values_.count(slot.value.value) != 0; + if (!written_value) { + continue; + } + Status status; + const ggml_tensor * tensor = resolve_external_slot(graph, slot, status); + if (tensor == nullptr) { + continue; + } + ggml_backend_hrx_imported_region region = {}; + if (!ggml_backend_hrx_get_retained_region(tensor, region)) { + continue; + } + const bool first_write = marked_identities.insert(region.identity).second; + if (first_write) { + ggml_backend_hrx_mark_retained_region_written(tensor); + } + } +} + PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { std::lock_guard lock(prepared_mutex_); return prepared_stats_; } +std::shared_ptr> GraphProgram::cached_weight_demand_for( + const CommandProgramBindings & bindings) const { + std::lock_guard lock(prepared_mutex_); + const bool prepared_current = has_prepared_ && resources_ != nullptr; + const bool cached_for_prepared = weight_demand_resources_.lock() == resources_; + const bool has_cached_demand = weight_demand_ != nullptr; + const bool reusable_demand = prepared_current && cached_for_prepared && has_cached_demand; + if (!reusable_demand) { + return nullptr; + } + size_t index = 0; + for (const CommandProgramBinding & binding : bindings.bindings()) { + if (weight_demand_values_.count(binding.value.value) == 0) { + continue; + } + const bool has_cached_binding = index < weight_demand_bindings_.size(); + const bool same_binding = has_cached_binding && + same_weight_demand_binding(binding, weight_demand_bindings_[index]); + if (!same_binding) { + return nullptr; + } + ++index; + } + if (index != weight_demand_bindings_.size()) { + return nullptr; + } + ++weight_demand_hits_; + return weight_demand_; +} + +void GraphProgram::remember_weight_demand(const CommandProgramBindings & bindings, + const std::vector & demand) { + std::vector candidates; + candidates.reserve(weight_demand_slots_.size()); + for (const CommandProgramBinding & binding : bindings.bindings()) { + if (weight_demand_values_.count(binding.value.value) != 0) { + candidates.push_back(binding); + } + } + std::lock_guard lock(prepared_mutex_); + const bool prepared_current = has_prepared_ && resources_ != nullptr; + if (!prepared_current) { + return; + } + weight_demand_bindings_ = std::move(candidates); + weight_demand_ = std::make_shared>(demand); + weight_demand_resources_ = resources_; + ++weight_demand_builds_; +} + +std::pair GraphProgram::weight_demand_stats() const { + std::lock_guard lock(prepared_mutex_); + return { weight_demand_builds_, weight_demand_hits_ }; +} + +void GraphProgram::invalidate_obsolete_weights(const std::vector & current_demand) { + std::lock_guard lock(prepared_mutex_); + if (!has_prepared_) { + return; + } + bool obsolete = false; + for (const HostWeightLease & lease : resources_->prepared.resident_host_weights) { + const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), + [&](const HostWeightSource & source) { return lease.matches(source); }); + obsolete |= !demanded; + } + if (obsolete) { + resources_.reset(); + has_prepared_ = false; + prepared_bindings_fingerprint_.clear(); + } +} + PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings) { + const CommandProgramBindings & bindings, + bool full_bindings) { PreparedCommandProgramCacheExecutionResult result; if (commands_ == nullptr || !commands_->valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; @@ -509,24 +743,53 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } std::lock_guard lock(prepared_mutex_); + const std::string bindings_fingerprint = + full_bindings ? command_program_bindings_fingerprint(bindings).value : std::string{}; + const bool fingerprint_changed = prepared_bindings_fingerprint_ != bindings_fingerprint; + const bool prepared_bindings_changed = full_bindings && fingerprint_changed; + const bool replace_prepared = has_prepared_ && prepared_bindings_changed; + if (replace_prepared) { + // A retained source may have changed while the graph structure stayed identical. Retire old bindings only + // after queued work completes; replay captures physical addresses that cannot be rebound in place. + if (ErrorResult error = take_status(hrx_stream_synchronize(context.stream))) { + result.status.log("synchronize before replacing HRX weight bindings failed: %s", error->c_str()); + return result; + } + context.graph_replay_state->mark_stream_synchronized(); + resources_.reset(); + has_prepared_ = false; + } if (!has_prepared_) { - prepared_ = prepare_command_program(context, *commands_, bindings); - if (!prepared_.valid()) { - result.status.append(prepared_.status); + resources_ = std::make_shared(); + resources_->prepared = prepare_command_program(context, *commands_, bindings); + if (!resources_->prepared.valid()) { + result.status.append(resources_->prepared.status); + resources_.reset(); return result; } has_prepared_ = true; + prepared_bindings_fingerprint_ = bindings_fingerprint; + for (ImportedSourceWatch & watch : imported_source_watches_) { + ggml_backend_hrx_imported_region region = {}; + if (ggml_backend_hrx_get_retained_region(watch.tensor, region)) { + watch.identity = region.identity; + watch.generation = region.generation; + watch.offset = region.offset; + watch.length = region.length; + } + } ++prepared_stats_.builds; } else { ++prepared_stats_.hits; } + result.execution_resources = resources_; if (environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_REPLAY") || debug_serial_command_execution_enabled()) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; result.graph_replay_ineligible_reason = debug_serial_command_execution_enabled() ? "debug_serial_execution" : "disabled_by_environment"; - result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, prepared_); + result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, resources_->prepared); if (!result.success) { result.status.log("execute cached HRX command program failed"); } @@ -534,11 +797,12 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } if (environment_flag_enabled("GGML_HRX_REBUILD_GRAPH_REPLAY")) { - recorded_ = {}; + resources_->recorded = {}; } const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, *commands_, bindings, prepared_, recorded_); + bind_and_launch_recorded_command_graph(context, *commands_, bindings, resources_->prepared, + resources_->recorded); apply_graph_replay_result(result, replay); if (replay.success) { result.success = true; @@ -551,7 +815,7 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } return result; } - result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, prepared_); + result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, resources_->prepared); if (!result.success) { result.status.log("execute cached HRX command program failed"); } @@ -650,7 +914,16 @@ GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_ if (can_skip_external_binding(*value, *commands_, values.size(), command_graph_values)) { continue; } - result.external_bindings.push_back({ id, tensor_by_value[static_cast(id.value)] }); + const ggml_tensor * tensor = tensor_by_value[static_cast(id.value)]; + const GraphProgramExternalSlot * slot = find_external_slot(id); + const bool has_slot = slot != nullptr; + const bool current_layout_eligible = ggml_backend_hrx_weight_layout_eligible(tensor); + const bool source_class_changed = has_slot && current_layout_eligible != slot->weight_layout_eligible; + if (source_class_changed) { + result.status.log("external weight source class changed for value %d", id.value); + return result; + } + result.external_bindings.push_back({ id, tensor }); } return result; } @@ -824,6 +1097,9 @@ GraphProgramCacheStats GraphProgramCache::stats() const { const PreparedCommandProgramCacheStats prepared = entry.second->prepared_stats(); stats.prepared_program_builds += prepared.builds; stats.prepared_program_hits += prepared.hits; + const auto [demand_builds, demand_hits] = entry.second->weight_demand_stats(); + stats.weight_demand_builds += demand_builds; + stats.weight_demand_hits += demand_hits; } return stats; } @@ -835,6 +1111,13 @@ void GraphProgramCache::clear() { last_program_ = nullptr; } +void GraphProgramCache::invalidate_obsolete_weights(const std::vector & current_demand) { + std::lock_guard lock(mutex_); + for (auto & program : programs_) { + program.second->invalidate_obsolete_weights(current_demand); + } +} + std::unique_ptr GraphProgramCache::build_program_from_imported(uint64_t uid, Graph && imported_graph, const KernelCorpus & corpus, diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index 4cdb3774370a..94e85f2cc642 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -3,6 +3,7 @@ #include "dispatch/command-program.h" #include "graph/graph.h" #include "kernel-corpus/kernel-corpus.h" +#include "runtime/command-program-executor.h" #include "runtime/prepared-command-program-cache.h" #include "status.h" @@ -11,6 +12,8 @@ #include #include #include +#include +#include #include struct ggml_cgraph; @@ -33,6 +36,7 @@ struct GraphProgramExternalSlot { GraphProgramExternalSlotKind kind = GraphProgramExternalSlotKind::Node; size_t node_index = 0; int32_t source_index = -1; + bool weight_layout_eligible = true; }; struct GraphProgramMatch { @@ -64,21 +68,32 @@ class GraphProgram { const CommandProgram & commands() const { return *commands_; } - CommandProgram & commands() { return *commands_; } + const CommandProgramWeightRequests & weight_requests() const { return weight_requests_; } GraphProgramMatch match_current_graph(const ggml_cgraph & graph) const; GraphProgramMatch match_trusted_graph(const ggml_cgraph & graph, bool bind_external = true) const; GraphProgramMatch match_host_staging_graph(const ggml_cgraph & graph) const; + GraphProgramMatch match_weight_demand_graph(const ggml_cgraph & graph) const; + std::shared_ptr> cached_weight_demand_for( + const CommandProgramBindings & bindings) const; + void remember_weight_demand(const CommandProgramBindings & bindings, + const std::vector & demand); Status capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match); bool has_prepared_program() const; bool can_use_prepared_fast_path(const ggml_cgraph & graph) const; + bool imported_sources_current() const; + void mark_imported_writes(const ggml_cgraph & graph) const; PreparedCommandProgramCacheStats prepared_stats() const; + std::pair weight_demand_stats() const; + + void invalidate_obsolete_weights(const std::vector & current_demand); PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings); + const CommandProgramBindings & bindings, + bool full_bindings); private: const GraphProgramExternalSlot * find_external_slot(ValueId value) const; @@ -89,17 +104,35 @@ class GraphProgram { uint64_t uid_ = 0; std::string target_; std::unique_ptr graph_; - std::unique_ptr commands_; + std::unique_ptr commands_; + const CommandProgramWeightRequests weight_requests_; std::string command_shape_; uint64_t command_shape_hash_ = 0; std::vector external_slots_; + std::vector weight_demand_slots_; + std::unordered_set weight_demand_values_; + struct ImportedSourceWatch { + const ggml_tensor * tensor; + uint64_t identity; + uint64_t generation; + size_t offset; + size_t length; + bool tracks_content_generation; + }; + std::vector imported_source_watches_; + std::unordered_set written_values_; std::unordered_map external_slot_by_value_; const ggml_tensor * const * fast_path_nodes_ = nullptr; mutable std::mutex prepared_mutex_; - PreparedCommandProgram prepared_; - RecordedCommandGraph recorded_; + std::shared_ptr resources_; + std::string prepared_bindings_fingerprint_; PreparedCommandProgramCacheStats prepared_stats_; + std::vector weight_demand_bindings_; + std::shared_ptr> weight_demand_; + std::weak_ptr weight_demand_resources_; + mutable uint64_t weight_demand_builds_ = 0; + mutable uint64_t weight_demand_hits_ = 0; bool has_prepared_ = false; }; @@ -108,6 +141,8 @@ struct GraphProgramCacheStats { uint64_t hits = 0; uint64_t prepared_program_builds = 0; uint64_t prepared_program_hits = 0; + uint64_t weight_demand_builds = 0; + uint64_t weight_demand_hits = 0; }; struct GraphProgramLookup { @@ -128,6 +163,7 @@ struct GraphProgramSupportResult { class GraphProgramCache { public: + void invalidate_obsolete_weights(const std::vector & current_demand); bool can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target) const; GraphProgramSupportResult check_support(const ggml_cgraph & graph, diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp index 02b591fcfc9c..fe9b04153f20 100644 --- a/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp @@ -7,7 +7,11 @@ namespace ggml::hrx { -HostBufferRef::HostBufferRef(hrx_buffer_t buffer, size_t offset) : buffer_(buffer), offset_(offset) {} +HostBufferRef::HostBufferRef(hrx_buffer_t buffer, size_t offset, size_t capacity, bool direct_bindable, + uintptr_t base, uint64_t identity, + std::shared_ptr> generation) : + buffer_(buffer), offset_(offset), capacity_(capacity), direct_bindable_(direct_bindable), + base_(base), identity_(identity), generation_(std::move(generation)) {} HostBufferRef::~HostBufferRef() { if (buffer_ != nullptr) { @@ -17,8 +21,13 @@ HostBufferRef::~HostBufferRef() { HostBufferRef::HostBufferRef(HostBufferRef && other) noexcept : buffer_(std::exchange(other.buffer_, nullptr)), - offset_(other.offset_) { + offset_(other.offset_), + capacity_(other.capacity_), + direct_bindable_(other.direct_bindable_), base_(other.base_), identity_(other.identity_), + generation_(std::move(other.generation_)) { other.offset_ = 0; + other.capacity_ = 0; + other.direct_bindable_ = false; } HostBufferRef & HostBufferRef::operator=(HostBufferRef && other) noexcept { @@ -28,17 +37,26 @@ HostBufferRef & HostBufferRef::operator=(HostBufferRef && other) noexcept { } buffer_ = std::exchange(other.buffer_, nullptr); offset_ = other.offset_; + capacity_ = other.capacity_; + direct_bindable_ = other.direct_bindable_; + base_ = other.base_; + identity_ = other.identity_; + generation_ = std::move(other.generation_); other.offset_ = 0; + other.capacity_ = 0; + other.direct_bindable_ = false; } return *this; } -void HostBufferRegistry::add(hrx_buffer_t buffer, void * base, size_t size) { +void HostBufferRegistry::add(hrx_buffer_t buffer, void * base, size_t size, bool direct_bindable, + uint64_t identity, std::shared_ptr> generation) { if (buffer == nullptr || base == nullptr || size == 0) { return; } std::lock_guard lock(mutex_); - entries_.push_back({ buffer, reinterpret_cast(base), size }); + entries_.push_back({ buffer, reinterpret_cast(base), size, direct_bindable, + identity, std::move(generation) }); } void HostBufferRegistry::remove(hrx_buffer_t buffer) { @@ -63,7 +81,8 @@ HostBufferRef HostBufferRegistry::find(const void * data, size_t size) const { continue; } hrx_buffer_retain(entry.buffer); - return HostBufferRef(entry.buffer, offset); + return HostBufferRef(entry.buffer, offset, entry.size, entry.direct_bindable, + entry.base, entry.identity, entry.generation); } return {}; } diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h index 30ce0eba2ec9..40971a87e020 100644 --- a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h @@ -1,7 +1,9 @@ #pragma once +#include #include #include +#include #include #include @@ -25,19 +27,31 @@ class HostBufferRef { hrx_buffer_t buffer() const { return buffer_; } size_t offset() const { return offset_; } + size_t capacity() const { return capacity_; } + bool direct_bindable() const { return direct_bindable_; } + void * base() const { return reinterpret_cast(base_); } + uint64_t identity() const { return identity_; } + std::shared_ptr> generation() const { return generation_; } private: hrx_buffer_t buffer_ = nullptr; size_t offset_ = 0; - - HostBufferRef(hrx_buffer_t buffer, size_t offset); + size_t capacity_ = 0; + bool direct_bindable_ = false; + uintptr_t base_ = 0; + uint64_t identity_ = 0; + std::shared_ptr> generation_; + + HostBufferRef(hrx_buffer_t buffer, size_t offset, size_t capacity, bool direct_bindable, + uintptr_t base, uint64_t identity, std::shared_ptr> generation); friend class HostBufferRegistry; }; // Tracks mapped HRX host buffers so pointer-based GGML transfers can remain stream ordered and handle based. class HostBufferRegistry { public: - void add(hrx_buffer_t buffer, void * base, size_t size); + void add(hrx_buffer_t buffer, void * base, size_t size, bool direct_bindable, + uint64_t identity, std::shared_ptr> generation); void remove(hrx_buffer_t buffer); HostBufferRef find(const void * data, size_t size) const; @@ -47,6 +61,9 @@ class HostBufferRegistry { hrx_buffer_t buffer = nullptr; uintptr_t base = 0; size_t size = 0; + bool direct_bindable = false; + uint64_t identity = 0; + std::shared_ptr> generation; }; mutable std::mutex mutex_; diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index c1741fd802ec..091d443951a8 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -1,6 +1,7 @@ #include "host-memory.h" #include "ggml-quants.h" +#include "ggml-impl.h" #include "hrx-interop-utils.h" #include "hrx_runtime.h" @@ -10,7 +11,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -21,6 +24,7 @@ namespace { static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; static constexpr size_t kLargeHostUploadBytes = 1024 * 1024; +static constexpr size_t kQ6PackScratchLimitBytes = 64 * 1024 * 1024; static bool mapped_download_staging_enabled() { static const bool enabled = [] { @@ -1427,16 +1431,39 @@ void HostTransferManager::clear() { stats_ = {}; } +struct HostWeightAccounting { + std::mutex mutex; + size_t allocation_count = 0; + size_t resident_bytes = 0; + size_t peak_resident_bytes = 0; +}; + struct HostWeightLease::Entry { ~Entry() { if (buffer != nullptr) { hrx_buffer_release(buffer); } + if (accounting != nullptr) { + std::lock_guard lock(accounting->mutex); + --accounting->allocation_count; + accounting->resident_bytes -= length; + } } hrx_buffer_t buffer = nullptr; size_t length = 0; std::string layout; + uint64_t source_identity = 0; + uint64_t source_generation = 0; + size_t source_capacity = 0; + size_t source_offset = 0; + size_t source_length = 0; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; + std::array source_ne = {}; + std::array source_nb = {}; + std::shared_ptr accounting; }; HostWeightLease::HostWeightLease(std::shared_ptr entry) : entry_(std::move(entry)) {} @@ -1458,6 +1485,22 @@ const std::string & HostWeightLease::layout() const { return entry_ != nullptr ? entry_->layout : empty; } +bool HostWeightLease::matches(const HostWeightSource & source) const { + if (!valid()) { + return false; + } + const Entry & entry = *entry_; + const size_t materialized_length = source.layout == kNativeWeightLayout ? source.length : source.materialized_length; + return entry.source_identity == source.identity && entry.source_generation == source.generation && + entry.source_capacity == source.capacity && entry.source_offset == source.offset && + entry.source_length == source.length && entry.length == materialized_length && + entry.layout == source.layout && entry.source_type == source.source_type && + entry.input_size == source.input_size && entry.output_size == source.output_size && + entry.source_ne == source.source_ne && entry.source_nb == source.source_nb; +} + +HostWeightCache::HostWeightCache() : accounting_(std::make_shared()) {} + HostWeightCache::~HostWeightCache() { clear(); } @@ -1480,6 +1523,10 @@ size_t HostWeightCache::SourceKeyHash::operator()(const SourceKey & key) const { mix(static_cast(key.source_type)); mix(static_cast(key.input_size)); mix(static_cast(key.output_size)); + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + mix(static_cast(key.source_ne[dimension])); + mix(static_cast(key.source_nb[dimension])); + } return static_cast(hash); } @@ -1513,15 +1560,21 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, source.length, materialized_length, source.layout, source.source_type, - source.input_size, source.output_size }; - { - std::lock_guard lock(mutex_); - const auto found = entries_.find(key); - if (found != entries_.end()) { - ++stats_.hits; - result.lease = HostWeightLease(found->second); - return result; - } + source.input_size, source.output_size, source.source_ne, source.source_nb }; + std::lock_guard lock(mutex_); + const auto found = entries_.find(key); + if (found != entries_.end()) { + ++stats_.hits; + result.lease = HostWeightLease(found->second); + return result; + } + if (fail_cold_after_ == 0) { + fail_cold_after_ = SIZE_MAX; + result.status.log("injected host weight acquisition failure"); + return result; + } + if (fail_cold_after_ != SIZE_MAX) { + --fail_cold_after_; } HostWeightSource materialization_source = source; @@ -1553,10 +1606,25 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device materialization_source.offset = 0; } + const bool q6_packed = source.layout == kQ6KPackedK256Row64ScaleRowLayout; + const bool scratch_fits = materialized_length <= kQ6PackScratchLimitBytes; + auto scratch = q6_pack_scratch_.find(materialized_length); + const bool scratch_exists = scratch != q6_pack_scratch_.end(); + const bool scratch_has_room = materialized_length <= kQ6PackScratchLimitBytes - q6_pack_scratch_bytes_; + const bool use_scratch = q6_packed && scratch_fits && (scratch_exists || scratch_has_room); + const bool create_scratch = use_scratch && !scratch_exists; + if (create_scratch) { + scratch = q6_pack_scratch_.emplace(materialized_length, std::vector{}).first; + q6_pack_scratch_bytes_ += materialized_length; + } + const void * upload_data = nullptr; size_t upload_size = 0; std::vector transformed; - result.status = materialize_weight(materialization_source, upload_data, upload_size, transformed); + // Budget 64 MiB of host scratch across phase transitions. The Q6 packer + // writes every output byte, and the cache lock covers packing and upload. + std::vector & output = use_scratch ? scratch->second : transformed; + result.status = materialize_weight(materialization_source, upload_data, upload_size, output); if (mapped_canonical.lock.owns_lock()) { mapped_canonical.lock.unlock(); } @@ -1572,40 +1640,148 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device auto entry = std::make_shared(); entry->length = materialized_length; entry->layout = source.layout; + entry->source_identity = source.identity; + entry->source_generation = source.generation; + entry->source_capacity = source.capacity; + entry->source_offset = source.offset; + entry->source_length = source.length; + entry->source_type = source.source_type; + entry->input_size = source.input_size; + entry->output_size = source.output_size; + entry->source_ne = source.source_ne; + entry->source_nb = source.source_nb; result.status = allocate_device_buffer(device, materialized_length, entry->buffer); if (!result.status.success()) { return result; } + entry->accounting = accounting_; + { + std::lock_guard accounting_lock(accounting_->mutex); + ++accounting_->allocation_count; + accounting_->resident_bytes += materialized_length; + accounting_->peak_resident_bytes = std::max(accounting_->peak_resident_bytes, accounting_->resident_bytes); + } result.status = transfers.upload_synchronous(stream, upload_data, entry->buffer, 0, upload_size); if (!result.status.success()) { return result; } - { - std::lock_guard lock(mutex_); - const auto inserted = entries_.emplace(key, entry); - if (!inserted.second) { - ++stats_.hits; - result.lease = HostWeightLease(inserted.first->second); - return result; - } - ++stats_.misses; - stats_.allocation_count = entries_.size(); - stats_.resident_bytes += materialized_length; - } + entries_.emplace(key, entry); + ++stats_.misses; + stats_.materialized_bytes += materialized_length; result.lease = HostWeightLease(std::move(entry)); return result; } HostWeightCacheStats HostWeightCache::stats() const { std::lock_guard lock(mutex_); - return stats_; + std::lock_guard accounting_lock(accounting_->mutex); + HostWeightCacheStats result = stats_; + result.allocation_count = accounting_->allocation_count; + result.resident_bytes = accounting_->resident_bytes; + result.peak_resident_bytes = accounting_->peak_resident_bytes; + return result; +} + +void HostWeightCache::report_coexisting_representations(const std::vector & leases, + const char * reason) { + struct Allocation { + hrx_buffer_t buffer = nullptr; + std::string layout; + size_t bytes = 0; + uint64_t generation = 0; + }; + struct Region { + uint64_t identity = 0; + size_t offset = 0; + size_t length = 0; + std::vector allocations; + }; + std::map regions; + for (const HostWeightLease & lease : leases) { + if (!lease.valid()) { + continue; + } + const HostWeightLease::Entry & entry = *lease.entry_; + std::ostringstream key; + key << entry.source_identity << ':' << entry.source_capacity << ':' + << entry.source_offset << ':' << entry.source_length << ':' << static_cast(entry.source_type); + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + key << ':' << entry.source_ne[dimension] << ':' << entry.source_nb[dimension]; + } + Region & region = regions[key.str()]; + region.identity = entry.source_identity; + region.offset = entry.source_offset; + region.length = entry.source_length; + const bool already_recorded = std::any_of(region.allocations.begin(), region.allocations.end(), + [&](const Allocation & allocation) { return allocation.buffer == entry.buffer; }); + if (!already_recorded) { + region.allocations.push_back({ entry.buffer, entry.layout, entry.length, entry.source_generation }); + } + } + std::lock_guard lock(mutex_); + for (const auto & [key, region] : regions) { + if (region.allocations.size() < 2) { + continue; + } + std::vector sorted = region.allocations; + std::sort(sorted.begin(), sorted.end(), [](const Allocation & lhs, const Allocation & rhs) { + return std::less{}(lhs.buffer, rhs.buffer); + }); + std::ostringstream signature; + signature << key; + std::ostringstream details; + size_t total_bytes = 0; + for (const Allocation & allocation : sorted) { + signature << ':' << allocation.buffer << ':' << allocation.layout << ':' << allocation.bytes + << ':' << allocation.generation; + details << (total_bytes == 0 ? "" : ", ") << allocation.layout << "@generation=" + << allocation.generation << "=" << allocation.bytes << "B"; + total_bytes += allocation.bytes; + } + if (!reported_duplicate_signatures_.insert(signature.str()).second) { + continue; + } + ++stats_.duplicate_warnings; + GGML_LOG_WARN("HRX weight residency: source=%llu region=[%zu,%zu) layouts={%s} resident_bytes=%zu reason=%s\n", + static_cast(region.identity), + region.offset, region.offset + region.length, details.str().c_str(), total_bytes, + reason == nullptr ? "unspecified" : reason); + } +} + +void HostWeightCache::fail_cold_acquire_after_for_test(size_t successful_acquires) { + std::lock_guard lock(mutex_); + fail_cold_after_ = successful_acquires; +} + +void HostWeightCache::evict_except(const std::vector & current_demand, + const std::vector & outstanding_leases) { + std::lock_guard lock(mutex_); + for (auto entry = entries_.begin(); entry != entries_.end();) { + const HostWeightLease resident(entry->second); + const bool current = std::any_of(current_demand.begin(), current_demand.end(), + [&](const HostWeightSource & source) { + return resident.matches(source); + }); + const bool outstanding = std::any_of(outstanding_leases.begin(), outstanding_leases.end(), + [&](const HostWeightLease & lease) { + return lease.entry_ == entry->second; + }); + const bool retained = current || outstanding; + if (retained) { + ++entry; + } else { + entry = entries_.erase(entry); + } + } } void HostWeightCache::clear() { std::lock_guard lock(mutex_); entries_.clear(); - stats_ = {}; + q6_pack_scratch_.clear(); + q6_pack_scratch_bytes_ = 0; } HostStagingBuffer::~HostStagingBuffer() { diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index a125598f20bf..cad5bb4da6f8 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -6,11 +6,14 @@ #include "status.h" #include +#include #include #include #include #include #include +#include +#include typedef struct hrx_buffer_s * hrx_buffer_t; typedef struct hrx_device_s * hrx_device_t; @@ -75,20 +78,30 @@ struct HostWeightSource { size_t capacity = 0; size_t offset = 0; size_t length = 0; + // Native extent and type identify the logical weight across its requested layouts. + size_t region_length = 0; + ggml_type region_type = GGML_TYPE_COUNT; size_t materialized_length = 0; std::string layout = kNativeWeightLayout; ggml_type source_type = GGML_TYPE_COUNT; int64_t input_size = 0; int64_t output_size = 0; + std::array source_ne = {}; + std::array source_nb = {}; }; struct HostWeightCacheStats { uint64_t hits = 0; uint64_t misses = 0; + size_t materialized_bytes = 0; size_t allocation_count = 0; size_t resident_bytes = 0; + size_t peak_resident_bytes = 0; + uint64_t duplicate_warnings = 0; }; +struct HostWeightAccounting; + class HostWeightLease { public: HostWeightLease() = default; @@ -97,6 +110,7 @@ class HostWeightLease { hrx_buffer_t buffer() const; size_t length() const; const std::string & layout() const; + bool matches(const HostWeightSource & source) const; private: struct Entry; @@ -115,7 +129,7 @@ struct HostWeightAcquireResult { class HostWeightCache { public: - HostWeightCache() = default; + HostWeightCache(); ~HostWeightCache(); HostWeightCache(const HostWeightCache &) = delete; @@ -126,6 +140,12 @@ class HostWeightCache { HostTransferManager & transfers, const HostWeightSource & source); HostWeightCacheStats stats() const; + // Report only distinct live device allocations for the same retained source region. + void report_coexisting_representations(const std::vector & leases, const char * reason); + void evict_except(const std::vector & current_demand, + const std::vector & outstanding_leases); + void fail_cold_acquire_after_for_test(size_t successful_acquires); + // Erases lookup ownership; leases and cumulative diagnostics outlive map entries. void clear(); private: @@ -140,13 +160,15 @@ class HostWeightCache { ggml_type source_type = GGML_TYPE_COUNT; int64_t input_size = 0; int64_t output_size = 0; + std::array source_ne = {}; + std::array source_nb = {}; bool operator==(const SourceKey & other) const { return identity == other.identity && generation == other.generation && capacity == other.capacity && offset == other.offset && length == other.length && materialized_length == other.materialized_length && layout == other.layout && source_type == other.source_type && input_size == other.input_size && - output_size == other.output_size; + output_size == other.output_size && source_ne == other.source_ne && source_nb == other.source_nb; } }; @@ -156,7 +178,13 @@ class HostWeightCache { mutable std::mutex mutex_; std::unordered_map, SourceKeyHash> entries_; + // Size-keyed scratch is overwritten on each use and is not a source-indexed representation cache. + std::unordered_map> q6_pack_scratch_; + size_t q6_pack_scratch_bytes_ = 0; HostWeightCacheStats stats_; + std::shared_ptr accounting_; + std::unordered_set reported_duplicate_signatures_; + size_t fail_cold_after_ = SIZE_MAX; }; struct HostStagingBuffer { diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index ad81b1a23afc..2968100f6d33 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -1,5 +1,6 @@ #include "prepared-command-program-cache.h" +#include #include #include #include @@ -98,12 +99,14 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ if (graph_uid == 0 || !commands.valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; result.graph_replay_ineligible_reason = "uncached_graph"; - PreparedCommandProgram prepared = prepare_command_program(context, commands, bindings); - if (!prepared.valid()) { - result.status.append(prepared.status); + auto resources = std::make_shared(); + resources->prepared = prepare_command_program(context, commands, bindings); + if (!resources->prepared.valid()) { + result.status.append(resources->prepared.status); return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, prepared); + result.execution_resources = resources; + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); if (!result.success) { result.status.log("execute uncached HRX command program failed"); } @@ -112,33 +115,33 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ const Key key = cache_key(graph_uid, context, command_shape_hash, bindings); std::shared_ptr entry; - bool created_entry = false; { std::lock_guard lock(mutex_); auto found = programs_.find(key); if (found == programs_.end()) { entry = std::make_shared(); programs_.emplace(key, entry); - created_entry = true; } else { entry = found->second; } } std::lock_guard entry_lock(entry->mutex); - if (entry->has_program && entry->program.valid()) { + if (entry->resources != nullptr && entry->resources->prepared.valid()) { + auto & resources = *entry->resources; + result.execution_resources = entry->resources; record_hit(); if (debug_serial_command_execution_enabled()) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; result.graph_replay_ineligible_reason = "debug_serial_execution"; - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources.prepared); if (!result.success) { result.status.log("execute cached HRX command program failed"); } return result; } const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); + bind_and_launch_recorded_command_graph(context, commands, bindings, resources.prepared, resources.recorded); apply_graph_replay_result(result, replay); if (replay.success) { result.success = true; @@ -151,33 +154,27 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ } return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources.prepared); if (!result.success) { result.status.log("execute cached HRX command program failed"); } return result; } - PreparedCommandProgram prepared = prepare_command_program(context, commands, bindings); - if (!prepared.valid()) { - if (created_entry) { - std::lock_guard lock(mutex_); - const auto found = programs_.find(key); - if (found != programs_.end() && found->second == entry && !entry->has_program) { - programs_.erase(found); - } - } - result.status.append(prepared.status); + auto resources = std::make_shared(); + resources->prepared = prepare_command_program(context, commands, bindings); + if (!resources->prepared.valid()) { + result.status.append(resources->prepared.status); return result; } - entry->program = std::move(prepared); - entry->has_program = true; + entry->resources = resources; + result.execution_resources = resources; record_build(); if (debug_serial_command_execution_enabled()) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; result.graph_replay_ineligible_reason = "debug_serial_execution"; - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); if (!result.success) { result.status.log("execute prepared HRX command program failed"); } @@ -185,7 +182,7 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ } const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); + bind_and_launch_recorded_command_graph(context, commands, bindings, resources->prepared, resources->recorded); apply_graph_replay_result(result, replay); if (replay.success) { result.success = true; @@ -198,7 +195,7 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ } return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); if (!result.success) { result.status.log("execute prepared HRX command program failed"); } @@ -215,6 +212,32 @@ void PreparedCommandProgramCache::clear() { programs_.clear(); } +void PreparedCommandProgramCache::invalidate_obsolete_weights(const std::vector & current_demand) { + std::vector> entries; + { + std::lock_guard lock(mutex_); + entries.reserve(programs_.size()); + for (const auto & program : programs_) { + entries.push_back(program.second); + } + } + for (const auto & entry : entries) { + std::lock_guard entry_lock(entry->mutex); + if (entry->resources == nullptr) { + continue; + } + bool obsolete = false; + for (const HostWeightLease & lease : entry->resources->prepared.resident_host_weights) { + const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), + [&](const HostWeightSource & source) { return lease.matches(source); }); + obsolete |= !demanded; + } + if (obsolete) { + entry->resources.reset(); + } + } +} + void PreparedCommandProgramCache::record_build() { std::lock_guard lock(mutex_); ++stats_.builds; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index 069ffd624480..911cb8c9de47 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -11,6 +11,7 @@ #include #include #include +#include namespace ggml::hrx { @@ -31,6 +32,7 @@ struct PreparedCommandProgramCacheExecutionResult { uint64_t graph_replay_total_ns = 0; size_t graph_replay_dispatches = 0; bool graph_replay_transient_allocation_changed = false; + std::shared_ptr execution_resources; }; class PreparedCommandProgramCache { @@ -56,6 +58,7 @@ class PreparedCommandProgramCache { PreparedCommandProgramCacheStats stats() const; void clear(); + void invalidate_obsolete_weights(const std::vector & current_demand); private: struct Key { @@ -81,9 +84,7 @@ class PreparedCommandProgramCache { struct Entry { std::mutex mutex; - PreparedCommandProgram program; - RecordedCommandGraph recorded; - bool has_program = false; + std::shared_ptr resources; }; void record_build(); diff --git a/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h b/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h new file mode 100644 index 000000000000..0573b2def49e --- /dev/null +++ b/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h @@ -0,0 +1,77 @@ +#pragma once + +#include "runtime/command-program-executor.h" +#include "runtime/host-memory.h" +#include "status.h" + +#include +#include +#include + +typedef struct hrx_event_s * hrx_event_t; +typedef struct hrx_semaphore_s * hrx_semaphore_t; + +namespace ggml::hrx { + +// Keeps submitted bindings alive until their stream timeline point or failure-path event completes. +// The current demand is retained separately so completed work can stay warm for the next call. +class WeightSubmissionTracker { + public: + ~WeightSubmissionTracker(); + + WeightSubmissionTracker(const WeightSubmissionTracker &) = delete; + WeightSubmissionTracker & operator=(const WeightSubmissionTracker &) = delete; + WeightSubmissionTracker() = default; + + Status poll_completed(bool * completed = nullptr); + Status record(hrx_device_t device, hrx_stream_t stream, + std::shared_ptr resources, + const std::vector & demand, bool admitted, bool demand_changed); + void mark_stream_synchronized(); + bool has_obsolete_pending() const; + void fail_completion_record_for_test() { fail_next_record_ = true; } + + const std::vector & current_demand() const { return current_demand_; } + bool demand_matches(const std::vector & demand) const; + std::vector merged_demand_for(const std::vector & demand) const; + void release_current_leases() { current_leases_.clear(); } + const std::vector & current_leases() const { return current_leases_; } + std::vector outstanding_leases() const; + + private: + struct RegionKey { + uint64_t identity = 0; + size_t offset = 0; + size_t length = 0; + ggml_type type = GGML_TYPE_COUNT; + std::array ne = {}; + std::array nb = {}; + + bool operator==(const RegionKey & other) const; + }; + struct RegionKeyHash { + size_t operator()(const RegionKey & key) const; + }; + using RegionDemands = std::unordered_map, RegionKeyHash>; + static RegionKey region_of(const HostWeightSource & source); + static RegionDemands group_demands(const std::vector & demand); + void publish_demand(const std::vector & demand); + + struct Pending { + hrx_event_t event = nullptr; + hrx_semaphore_t semaphore = nullptr; + uint64_t value = 0; + std::shared_ptr resources; + bool may_hold_obsolete_demand = false; + }; + + static void release_completion(Pending & pending); + + std::vector current_demand_; + RegionDemands retained_regions_; + std::vector current_leases_; + std::vector pending_; + bool fail_next_record_ = false; +}; + +} // namespace ggml::hrx diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 659cff9bb0af..655b12ec17de 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1028,13 +1028,18 @@ static void run_graph_import_checks() { REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_identity_bindings).value != runtime_fingerprint.value); - ggml::hrx::ValueMap changed_generation = imported.graph.values(); - REQUIRE(changed_generation.bind_buffer( - a_value->id, { dummy_hrx_buffer(0x1000), 0, a_value->byte_count, 0, 1, a_value->byte_count })); + REQUIRE(!a_binding->weight); + std::vector changed_generation_values = runtime_bindings.bindings(); + auto changed_input = std::find_if(changed_generation_values.begin(), changed_generation_values.end(), + [&](const auto & binding) { return binding.value == a_value->id; }); + REQUIRE(changed_input != changed_generation_values.end()); + ++changed_input->generation; const ggml::hrx::CommandProgramBindings changed_generation_bindings = - ggml::hrx::CommandProgramBindings::from_value_map(changed_generation); + ggml::hrx::CommandProgramBindings::from_bindings(std::move(changed_generation_values)); REQUIRE(changed_generation_bindings.valid()); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_generation_bindings).value != + REQUIRE(ggml::hrx::command_program_bindings_hash(changed_generation_bindings).value == + ggml::hrx::command_program_bindings_hash(runtime_bindings).value); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_generation_bindings).value == runtime_fingerprint.value); ggml::hrx::ValueMap changed_capacity = imported.graph.values(); @@ -1152,6 +1157,9 @@ static void run_graph_import_checks() { host_value_binding.generation = 1; host_value_binding.capacity = host_storage.size(); host_value_binding.weight = true; + host_value_binding.source_host_data = host_storage.data(); + host_value_binding.source_offset = 16; + host_value_binding.source_capacity = host_storage.size(); ggml::hrx::ValueMap host_values = imported.graph.values(); REQUIRE(host_values.bind_buffer(a_value->id, host_value_binding)); REQUIRE(host_values.bind_buffer(out_value->id, { dummy_hrx_buffer(0x5000), 0, out_value->byte_count })); @@ -1165,7 +1173,28 @@ static void run_graph_import_checks() { REQUIRE(host_binding->offset == 16); REQUIRE(host_binding->length == a_value->byte_count); REQUIRE(host_binding->weight); + ggml::hrx::ValueBufferBinding staging_value_binding = host_value_binding; + staging_value_binding.source_host_data = nullptr; + REQUIRE(host_values.bind_buffer(a_value->id, staging_value_binding)); + const auto staging_bindings = ggml::hrx::CommandProgramBindings::from_value_map(host_values); + REQUIRE(staging_bindings.valid()); + const auto * staging_binding = staging_bindings.find(a_value->id); + REQUIRE(staging_binding != nullptr); + REQUIRE(staging_binding->host_data == host_storage.data()); + REQUIRE(staging_binding->source_host_data == nullptr); REQUIRE(ggml::hrx::command_program_bindings_fingerprint(host_bindings).value != runtime_fingerprint.value); + std::vector changed_weight_values = host_bindings.bindings(); + auto changed_weight = std::find_if(changed_weight_values.begin(), changed_weight_values.end(), + [&](const auto & binding) { return binding.value == a_value->id; }); + REQUIRE(changed_weight != changed_weight_values.end()); + ++changed_weight->generation; + const auto changed_weight_bindings = ggml::hrx::CommandProgramBindings::from_bindings( + std::move(changed_weight_values)); + REQUIRE(changed_weight_bindings.valid()); + REQUIRE(ggml::hrx::command_program_bindings_hash(changed_weight_bindings).value != + ggml::hrx::command_program_bindings_hash(host_bindings).value); + REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_weight_bindings).value != + ggml::hrx::command_program_bindings_fingerprint(host_bindings).value); resolved = ggml::hrx::resolve_command_program_bindings(commands, missing_bindings); REQUIRE(!resolved.valid()); @@ -6163,6 +6192,8 @@ static void run_generic_swiglu_k16_publication_edge_checks() { }); REQUIRE(producer != plan.dispatches.end()); REQUIRE(producer->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(producer->kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(producer->kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); REQUIRE(producer->bindings.size() == 4); REQUIRE(producer->bindings.back().value == glu_value->id); @@ -6220,6 +6251,8 @@ static void run_generic_swiglu_k16_publication_edge_checks() { }); REQUIRE(producer != plan.dispatches.end()); REQUIRE(producer->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(producer->kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(producer->kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); REQUIRE(producer->bindings.size() == 4); const auto copy = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major" && @@ -6300,6 +6333,8 @@ static void run_tiled_matmul_alternate_publish_checks() { } else { REQUIRE(producer->bindings.size() == 3); REQUIRE(producer->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(producer->kernel.integer_parameters.at("token_count") == 256); + REQUIRE(producer->kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); const auto quantize = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && @@ -6372,8 +6407,10 @@ static void run_tiled_matmul_alternate_publish_checks() { dispatch.bindings[2].value == projection_value->id; }); REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 3); REQUIRE(producer->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(producer->kernel.integer_parameters.at("token_count") == 256); + REQUIRE(producer->kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); + REQUIRE(producer->bindings.size() == 3); const auto q5_consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { @@ -6729,6 +6766,7 @@ static void run_rmsnorm_binary_k16_output_checks() { REQUIRE(diagnostics.unsupported_node != nullptr); REQUIRE(diagnostics.unsupported_node_index == 0); REQUIRE(diagnostics.unsupported_node->op == GGML_OP_RMS_NORM); + REQUIRE(!diagnostics.match.attempts.empty()); REQUIRE(!ggml::hrx::DispatchScheduler::supports_node( imported.graph, diagnostics.unsupported_node, test_dispatch_target())); REQUIRE(!ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target())); @@ -11445,16 +11483,24 @@ static void run_zero_output_scheduler_elision_checks() { REQUIRE(ids_value != nullptr); ggml::hrx::Command referenced; referenced.bindings.push_back({ "ids", ids_value->id }); - for (auto * list : - { &lookup.program->commands().initialization_commands, &lookup.program->commands().commands }) { - list->push_back(referenced); - const auto match = lookup.program->match_current_graph(*graph); + auto make_program = [&](ggml::hrx::CommandProgram commands) { + return std::make_unique( + lookup.program->uid(), lookup.program->target(), + std::make_unique(lookup.program->graph()), + std::make_unique(std::move(commands)), std::string{}); + }; + for (const bool initialization : { true, false }) { + auto commands = copy_command_program_shape(lookup.program->commands()); + auto & list = initialization ? commands.initialization_commands : commands.commands; + list.push_back(referenced); + const auto referenced_program = make_program(std::move(commands)); + const auto match = referenced_program->match_current_graph(*graph); REQUIRE(match.valid()); REQUIRE(match.external_bindings.size() == 2); REQUIRE(std::any_of(match.external_bindings.begin(), match.external_bindings.end(), [&](const auto & binding) { return binding.tensor == ids; })); - list->clear(); - const auto unreferenced = lookup.program->match_current_graph(*graph); + const auto unreferenced_program = make_program(copy_command_program_shape(lookup.program->commands())); + const auto unreferenced = unreferenced_program->match_current_graph(*graph); REQUIRE(unreferenced.valid()); REQUIRE(unreferenced.external_bindings.size() == 1); for (const auto origin : @@ -11463,12 +11509,15 @@ static void run_zero_output_scheduler_elision_checks() { unrelated.bindings.push_back({ "negative", ggml::hrx::ValueId(-1) }); unrelated.bindings.push_back({ "large", ggml::hrx::ValueId(INT32_MAX) }); unrelated.bindings.push_back({ "ids", ids_value->id, origin }); - list->push_back(unrelated); - const auto ignored = lookup.program->match_current_graph(*graph); + auto unrelated_commands = copy_command_program_shape(lookup.program->commands()); + auto & unrelated_list = + initialization ? unrelated_commands.initialization_commands : unrelated_commands.commands; + unrelated_list.push_back(unrelated); + const auto unrelated_program = make_program(std::move(unrelated_commands)); + const auto ignored = unrelated_program->match_current_graph(*graph); REQUIRE(ignored.valid()); REQUIRE(ignored.external_bindings.size() == 1); REQUIRE(ignored.external_bindings[0].tensor == input); - list->clear(); } } } @@ -11626,6 +11675,142 @@ static void run_zero_output_device_support_checks() { ggml_backend_free(backend); } +static void run_zero_output_unused_weight_staging(bool direct_execution) { + EnvironmentVariableGuard replay_mode("GGML_HRX_DISABLE_GRAPH_REPLAY"); + if (direct_execution) { + replay_mode.set("1"); + } else { + replay_mode.unset(); + } + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * state = static_cast(backend->context); + ggml_context * ctx = ggml_init({ 512 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + + constexpr int64_t width = 256; + constexpr int64_t unused_rows = 4096; + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, width); + ggml_tensor * unused_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, unused_rows); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, 1); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_set_output(output); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 0); + ggml_tensor * empty_rows = ggml_get_rows(ctx, output, ids); + ggml_tensor * empty_logits = ggml_mul_mat(ctx, unused_weight, empty_rows); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, empty_logits); + graph->uid = 3004; + + std::vector weights(width * (width + unused_rows), 0.0f); + for (int64_t i = 0; i < width; ++i) { + weights[i * width + i] = 1.0f; + } + const size_t weight_bytes = ggml_nbytes(weight); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), weights.data(), weights.size() * sizeof(float), ggml_nbytes(unused_weight)); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, weight, weights.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, unused_weight, weights.data() + width * width) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t host = ggml_backend_alloc_ctx_tensors_from_buft(ctx, ggml_backend_cpu_buffer_type()); + REQUIRE(host != nullptr); + + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(cpu != nullptr); + ggml_context * cpu_ctx = ggml_init({ 128 * 1024, nullptr, true }); + REQUIRE(cpu_ctx != nullptr); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, width, width); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, width, 1); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + REQUIRE(cpu_buffer != nullptr); + ggml_backend_tensor_set(cpu_weight, weights.data(), 0, weight_bytes); + + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + auto lookup = state->graph_programs.get_or_build(*graph, corpus, state->device->architecture); + for (const auto & error : lookup.status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + for (const auto & error : lookup.match.status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + REQUIRE(lookup.valid()); + const auto * unused_value = lookup.program->graph().values().find_tensor(unused_weight); + REQUIRE(unused_value != nullptr); + std::vector external_bindings; + for (const auto & external : lookup.match.external_bindings) { + ggml::hrx::ValueBufferBinding buffer; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); + const auto * value = lookup.program->graph().values().find(external.value); + REQUIRE(value != nullptr); + external_bindings.push_back(ggml::hrx::make_command_program_binding( + external.value, buffer, value->type, value->ne, value->nb, value->byte_count == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); + REQUIRE(bindings.valid()); + REQUIRE(bindings.find(unused_value->id) != nullptr); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + state->device->device, state->stream, state->device->architecture.c_str(), &corpus, + &state->kernel_executables, &state->transient_arena, &state->host_transfers, + &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, + }; + for (int iteration = 0; iteration < 2; ++iteration) { + std::vector values(width); + for (int64_t i = 0; i < width; ++i) { + values[i] = static_cast(i + iteration) * 0.25f; + } + const std::vector sentinel(width, -123.0f); + ggml_backend_tensor_set(output, sentinel.data(), 0, ggml_nbytes(output)); + ggml_backend_tensor_set(input, values.data(), 0, ggml_nbytes(input)); + ggml_backend_tensor_set(cpu_input, values.data(), 0, ggml_nbytes(cpu_input)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + std::vector expected(width); + ggml_backend_tensor_get(cpu_output, expected.data(), 0, ggml_nbytes(cpu_output)); + auto execution = lookup.program->execute_with_result(execution_context, bindings, true); + REQUIRE(execution.success); + ggml_backend_synchronize(backend); + const auto & prepared = execution.execution_resources->prepared; + REQUIRE(std::none_of(prepared.host_staging.begin(), prepared.host_staging.end(), [&](const auto & staging) { + return staging.value == unused_value->id.value; + })); + REQUIRE(prepared.host_staging.size() == 2); + REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(state->host_weights.stats().resident_bytes == weight_bytes); + REQUIRE(state->host_weights.stats().misses == 1); + // The vector matmul declares both host input and output as reads. + const size_t staging_upload_bytes = ggml_nbytes(input) + ggml_nbytes(output); + REQUIRE(state->host_transfers.stats().upload_bytes == weight_bytes + (iteration + 1) * staging_upload_bytes); + REQUIRE(lookup.program->prepared_stats().builds == 1); + REQUIRE(lookup.program->prepared_stats().hits == static_cast(iteration)); + const auto expected_event = direct_execution ? ggml::hrx::HrxGraphReplayEvent::Disabled : + (iteration == 0 ? ggml::hrx::HrxGraphReplayEvent::MissBuild : ggml::hrx::HrxGraphReplayEvent::Hit); + REQUIRE(execution.graph_replay_event == expected_event); + std::vector actual(width); + ggml_backend_tensor_get(output, actual.data(), 0, ggml_nbytes(output)); + REQUIRE(actual == expected); + } + // Initialization-only bindings need the same staging and weight access analysis. + auto initialization_only = copy_command_program_shape(lookup.program->commands()); + initialization_only.initialization_commands = std::move(initialization_only.commands); + { + auto prepared = ggml::hrx::prepare_command_program(execution_context, initialization_only, bindings); + REQUIRE(prepared.valid()); + REQUIRE(prepared.host_staging.size() == 2); + REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(state->host_weights.stats().misses == 1); + } + ggml_backend_buffer_free(cpu_buffer); + ggml_free(cpu_ctx); + ggml_backend_free(cpu); + ggml_backend_free(backend); + ggml_backend_buffer_free(host); + ggml_backend_buffer_free(imported); + ggml_free(ctx); +} + static void run_scale_f32_device_support_checks() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -12122,6 +12307,253 @@ static void run_disjoint_transient_plan_packing_checks() { ggml_free(ctx); } +static void run_command_weight_request_checks() { + using namespace ggml::hrx; + std::array source = {}; + CommandProgramBinding binding; + binding.value = ValueId(1); + binding.buffer = dummy_hrx_buffer(1); + binding.length = 64; + binding.identity = 10; + binding.generation = 20; + binding.weight = true; + binding.source_host_data = source.data(); + binding.source_capacity = source.size(); + binding.source_type = GGML_TYPE_F32; + binding.source_ne = { 8, 2, 1, 1 }; + binding.source_nb = { 4, 32, 64, 64 }; + + CommandBinding native = { "first-native", ValueId(1) }; + CommandBinding packed = native; + packed.name = "first-packed"; + packed.layout = "test-packed"; + packed.source_type = GGML_TYPE_F32; + packed.input_size = 8; + packed.output_size = 2; + packed.source_length = 64; + packed.length = 32; + CommandBinding distinct = packed; + distinct.name = "distinct-transform"; + distinct.source_length = 32; + distinct.input_size = 4; + CommandBinding init_only = native; + init_only.value = ValueId(2); + CommandBinding read_then_write = native; + read_then_write.value = ValueId(3); + CommandBinding read_write = native; + read_write.value = ValueId(4); + CommandBinding write_then_read = native; + write_then_read.value = ValueId(5); + CommandBinding initial_write = write_then_read; + initial_write.access = ResourceAccess::Write; + Command init; + init.bindings = { native, packed, init_only, read_then_write, read_write, initial_write }; + Command main; + CommandBinding native_duplicate = native; + native_duplicate.name = "later-native"; + native_duplicate.offset = 8; + CommandBinding packed_duplicate = packed; + packed_duplicate.name = "later-packed"; + read_then_write.access = ResourceAccess::Write; + read_write.access = ResourceAccess::ReadWrite; + CommandBinding transient_write = native; + transient_write.origin = CommandBindingOrigin::Transient; + transient_write.access = ResourceAccess::Write; + CommandBinding constant_read = packed; + constant_read.origin = CommandBindingOrigin::ProgramConstant; + main.bindings = { packed_duplicate, native_duplicate, distinct, read_then_write, read_write, + write_then_read, transient_write, constant_read }; + CommandProgram commands; + commands.initialization_commands.push_back(init); + commands.commands.push_back(main); + const auto requests = build_command_program_weight_requests(commands); + std::vector current; + for (int32_t value = 1; value <= 5; ++value) { + auto next = binding; + next.value = ValueId(value); + current.push_back(next); + } + const auto bindings = CommandProgramBindings::from_bindings(current); + REQUIRE(bindings.valid()); + const auto demands = bind_command_program_weight_demands(requests, bindings); + REQUIRE(demands.valid()); + REQUIRE(demands.weights.size() == 4); + REQUIRE(demands.weights[0].request.name == "first-native"); + REQUIRE(demands.weights[1].request.name == "first-packed"); + REQUIRE(demands.weights[2].request.value == ValueId(2)); + REQUIRE(demands.weights[3].request.name == "distinct-transform"); + REQUIRE(!demands.weights[0].requires_materialization); + REQUIRE(demands.weights[1].requires_materialization); + REQUIRE(demands.weights[1].source.materialized_length == 32); + REQUIRE(demands.weights[1].source.length == 64); + REQUIRE(demands.weights[3].source.length == 32); + REQUIRE(demands.weights[3].source.input_size == 4); + const auto collected = collect_command_program_weight_demands(commands, bindings); + REQUIRE(collected.valid()); + REQUIRE(collected.weights.size() == demands.weights.size()); + REQUIRE(collected.weights[2].request.value == ValueId(2)); + + // Reuse the immutable plan with different live source metadata and allocation eligibility. + commands.initialization_commands.clear(); + commands.commands.clear(); + for (auto & next : current) { + next.buffer = nullptr; + next.host_data = source.data() + 32; + next.identity = 30; + next.generation = 40; + next.length = 128; + next.source_host_data = source.data() + 16; + next.source_offset = 32; + next.source_capacity = source.size() - 16; + next.source_type = GGML_TYPE_F16; + next.source_ne = { 16, 4, 1, 1 }; + next.source_nb = { 2, 32, 128, 128 }; + } + const auto rebound = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings(current)); + REQUIRE(rebound.valid()); + REQUIRE(rebound.weights.size() == 4); + for (const auto & demand : rebound.weights) { + REQUIRE(demand.source.host_data == source.data() + 16); + REQUIRE(demand.source.identity == 30); + REQUIRE(demand.source.generation == 40); + REQUIRE(demand.source.offset == 32); + REQUIRE(demand.source.capacity == source.size() - 16); + REQUIRE(demand.source.region_length == 128); + REQUIRE(demand.source.region_type == GGML_TYPE_F16); + REQUIRE(demand.source.source_ne == current[0].source_ne); + REQUIRE(demand.source.source_nb == current[0].source_nb); + REQUIRE(demand.requires_materialization); + } + REQUIRE(rebound.weights[0].source.length == 128); + REQUIRE(rebound.weights[0].source.materialized_length == 128); + REQUIRE(rebound.weights[0].source.source_type == GGML_TYPE_F16); + REQUIRE(rebound.weights[1].source.length == 64); + REQUIRE(rebound.weights[1].source.source_type == GGML_TYPE_F32); + for (auto & next : current) { + next.weight = false; + next.source_host_data = nullptr; + } + const auto nonweights = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings(current)); + REQUIRE(nonweights.valid()); + REQUIRE(nonweights.weights.empty()); + + CommandProgram native_program; + native_program.initialization_commands.push_back(Command{}); + native_program.initialization_commands.back().bindings.push_back(native); + const auto native_requests = build_command_program_weight_requests(native_program); + auto uploaded = binding; + uploaded.source_host_data = nullptr; + const auto uploaded_demand = bind_command_program_weight_demands( + native_requests, CommandProgramBindings::from_bindings({ uploaded })); + REQUIRE(uploaded_demand.valid()); + REQUIRE(uploaded_demand.weights.empty()); + const auto restored = bind_command_program_weight_demands(native_requests, bindings); + REQUIRE(restored.valid()); + REQUIRE(restored.weights.size() == 1); + REQUIRE(!restored.weights[0].requires_materialization); + auto later_binding = binding; + later_binding.identity = 99; + const auto first_binding = bind_command_program_weight_demands( + native_requests, CommandProgramBindings::from_bindings({ binding, later_binding })); + REQUIRE(first_binding.valid()); + REQUIRE(first_binding.weights.size() == 1); + REQUIRE(first_binding.weights[0].source.identity == binding.identity); +} + +static void run_command_weight_request_validation_checks() { + using namespace ggml::hrx; + std::array source = {}; + CommandProgramBinding binding; + binding.value = ValueId(1); + binding.buffer = dummy_hrx_buffer(1); + binding.length = source.size(); + binding.weight = true; + binding.source_host_data = source.data(); + binding.source_capacity = source.size(); + CommandBinding valid = { "valid", ValueId(1) }; + valid.layout = "test-packed"; + valid.source_type = GGML_TYPE_F32; + valid.input_size = 8; + valid.output_size = 2; + valid.source_length = 64; + valid.length = 32; + CommandBinding invalid_offset = valid; + invalid_offset.name = "invalid-offset"; + invalid_offset.offset = 1; + for (const bool invalid_first : { false, true }) { + CommandProgram commands; + Command init; + init.bindings.push_back(invalid_first ? invalid_offset : valid); + commands.initialization_commands.push_back(init); + Command main; + main.bindings.push_back(invalid_first ? valid : invalid_offset); + commands.commands.push_back(main); + const auto requests = build_command_program_weight_requests(commands); + const auto demands = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings({ binding })); + REQUIRE(!demands.valid()); + REQUIRE(demands.status.errors().size() == 1); + REQUIRE(demands.status.errors()[0].find("invalid transformed layout request") != std::string::npos); + REQUIRE(demands.weights.size() == 1); + REQUIRE(demands.weights[0].request.name == "valid"); + REQUIRE(demands.weights[0].request.offset == 0); + auto missing = binding; + missing.source_host_data = nullptr; + const auto missing_source = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings({ missing })); + REQUIRE(!missing_source.valid()); + REQUIRE(missing_source.weights.empty()); + REQUIRE(missing_source.status.errors().size() == 2); + for (const auto & error : missing_source.status.errors()) { + REQUIRE(error.find("no retained source") != std::string::npos); + } + missing.weight = false; + const auto nonweight = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings({ missing })); + REQUIRE(nonweight.valid()); + REQUIRE(nonweight.weights.empty()); + const auto unbound = bind_command_program_weight_demands(requests, CommandProgramBindings{}); + REQUIRE(unbound.valid()); + REQUIRE(unbound.weights.empty()); + } + std::vector invalid(6, valid); + invalid[0].source_type = GGML_TYPE_COUNT; + invalid[1].input_size = 0; + invalid[2].output_size = 0; + invalid[3].source_length = 0; + invalid[4].source_length = source.size() + 1; + invalid[5].length = 0; + for (const auto & request : invalid) { + CommandProgram commands; + Command init; + init.bindings.push_back(request); + commands.initialization_commands.push_back(init); + const auto demands = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ binding })); + REQUIRE(!demands.valid()); + REQUIRE(demands.weights.empty()); + REQUIRE(demands.status.errors().size() == 1); + REQUIRE(demands.status.errors()[0].find("invalid transformed layout request") != std::string::npos); + } + CommandProgram commands; + Command init; + init.bindings.push_back(valid); + commands.initialization_commands.push_back(init); + const auto requests = build_command_program_weight_requests(commands); + auto short_binding = binding; + short_binding.length = valid.source_length - 1; + const auto too_short = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings({ short_binding })); + REQUIRE(!too_short.valid()); + REQUIRE(too_short.weights.empty()); + const auto restored = bind_command_program_weight_demands( + requests, CommandProgramBindings::from_bindings({ binding })); + REQUIRE(restored.valid()); + REQUIRE(restored.weights.size() == 1); +} + static void run_command_shape_hash_checks() { REQUIRE(ggml::hrx::command_program_shape_hash("") == UINT64_C(1469598103934665603)); std::string bytes; @@ -12130,7 +12562,8 @@ static void run_command_shape_hash_checks() { } bytes.append("\0key\xff", 5); REQUIRE(ggml::hrx::command_program_shape_hash(bytes) == UINT64_C(6643458358697029213)); - ggml::hrx::GraphProgram program(1, "gfx1151", nullptr, nullptr, bytes); + ggml::hrx::GraphProgram program(1, "gfx1151", std::make_unique(), + std::make_unique(), bytes); bytes.clear(); REQUIRE(program.command_shape_hash() == UINT64_C(6643458358697029213)); REQUIRE(program.command_shape_hash() == ggml::hrx::command_program_shape_hash(program.command_shape())); @@ -14135,6 +14568,8 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("multiple_transient_plan", [] { run_multiple_transient_plan_checks(); }); suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); + suite.host_case("command_weight_requests", [] { run_command_weight_request_checks(); }); + suite.host_case("command_weight_request_validation", [] { run_command_weight_request_validation_checks(); }); suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); suite.host_case("graph_compute_mask", [] { run_graph_compute_mask_checks(); }); suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); @@ -14218,6 +14653,8 @@ static void register_hrx_backend_device_cases(test_runner::Suite & suite) { suite.device_case("initializer_host_input_ownership", [] { run_initializer_host_input_ownership(); }); suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); + suite.device_case("zero_output_unused_weight_staging.recorded", [] { run_zero_output_unused_weight_staging(false); }); + suite.device_case("zero_output_unused_weight_staging.direct", [] { run_zero_output_unused_weight_staging(true); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); suite.device_case("binary_q8_publication_numerics", [] { run_binary_q8_publication_numerics(); }); suite.device_case("qwen_expert_table_partition_prefill_512_execution", diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index e281a3199f09..a4f782753394 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -2,18 +2,25 @@ #include "backend-context.h" #include "ggml-backend-impl.h" #include "ggml-backend.h" +#include "ggml-cpu.h" #include "ggml-hrx.h" +#include "ggml-impl.h" #include "ggml.h" #include "hrx-interop-utils.h" #include "runtime/host-memory.h" +#include "runtime/command-program-executor.h" #include "testing_suite.h" +#include #include #include #include #include #include #include +#if defined(__linux__) +#include +#endif #define REQUIRE(condition) \ do { \ @@ -162,6 +169,24 @@ static void run_host_buffer_checks(ggml_backend_t backend) { REQUIRE(direct_binding.host_data == nullptr); REQUIRE(direct_binding.offset == 0); REQUIRE(direct_binding.length == ggml_nbytes(direct_tensor)); + ggml_backend_hrx_imported_region retained_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(host_tensor, retained_region)); + REQUIRE(!ggml_backend_hrx_get_imported_region(host_tensor, retained_region)); + const uint64_t before_host_set = retained_region.generation; + const uint32_t host_replacement = 0xabcdef01; + ggml_backend_tensor_set(host_tensor, &host_replacement, 0, sizeof(host_replacement)); + REQUIRE(ggml_backend_hrx_get_retained_region(host_tensor, retained_region)); + REQUIRE(retained_region.generation != before_host_set); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + const uint64_t before_direct_set = retained_region.generation; + ggml_backend_tensor_set_async(backend, direct_tensor, &host_replacement, 0, sizeof(host_replacement)); + ggml_backend_synchronize(backend); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + REQUIRE(retained_region.generation != before_direct_set); + const uint64_t before_host_clear = retained_region.generation; + ggml_backend_buffer_clear(direct_buffer, 0); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + REQUIRE(retained_region.generation != before_host_clear); ggml_tensor * tensor = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 64); ggml_backend_buffer_t local = ggml_backend_alloc_buffer(backend, 4096); @@ -187,12 +212,231 @@ static void run_host_buffer_checks(ggml_backend_t backend) { REQUIRE(context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed) == upload_fallbacks); REQUIRE(context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed) == download_fallbacks); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + const uint64_t before_host_copy = retained_region.generation; + ggml_backend_tensor_copy_async(backend, backend, tensor, direct_tensor); + ggml_backend_synchronize(backend); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + REQUIRE(retained_region.generation != before_host_copy); + const uint64_t before_host_memset = retained_region.generation; + ggml_backend_tensor_memset(direct_tensor, 0x5a, 0, sizeof(uint32_t)); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + REQUIRE(retained_region.generation != before_host_memset); + ggml_backend_buffer_free(local); ggml_backend_buffer_free(direct_buffer); ggml_backend_buffer_free(buffer); ggml_free(ggml); } +static void run_imported_buffer_checks(ggml_backend_t backend) { + ggml_backend_dev_t device = ggml_backend_get_device(backend); + ggml_backend_dev_props props = {}; + ggml_backend_dev_get_props(device, &props); + REQUIRE(props.caps.buffer_from_host_ptr); + + std::array source = {}; + for (size_t i = 0; i < source.size(); ++i) { + source[i] = static_cast(i + 1); + } + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr(device, source.data(), sizeof(source), + sizeof(source) / 2); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_buffer_get_base(imported) == source.data()); + REQUIRE(!ggml_backend_buffer_is_host(imported)); + auto * context = ggml_backend_hrx_buffer_context_from_buffer(imported); + REQUIRE(context->imported); + REQUIRE(context->buffer == nullptr); + + ggml_init_params params = {}; + params.mem_size = 4096; + params.no_alloc = true; + ggml_context * ggml = ggml_init(params); + REQUIRE(ggml != nullptr); + ggml_tensor * first = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 32); + ggml_tensor * second = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 32); + REQUIRE(ggml_backend_tensor_alloc(imported, first, source.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, second, source.data() + 32) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + ggml_backend_hrx_imported_region first_region = {}; + ggml_backend_hrx_imported_region second_region = {}; + REQUIRE(ggml_backend_hrx_get_imported_region(first, first_region)); + REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); + REQUIRE(first_region.identity == second_region.identity); + REQUIRE(first_region.offset == 0); + REQUIRE(second_region.offset == 32 * sizeof(uint32_t)); + REQUIRE(first_region.type == GGML_TYPE_I32); + REQUIRE(first_region.ne[0] == 32); + REQUIRE(first_region.nb[0] == sizeof(uint32_t)); + + ggml_tensor * view = ggml_view_1d(ggml, second, 8, 4 * sizeof(uint32_t)); + REQUIRE(ggml_backend_view_init(view) == GGML_STATUS_SUCCESS); + ggml_backend_hrx_imported_region view_region = {}; + REQUIRE(ggml_backend_hrx_get_imported_region(view, view_region)); + REQUIRE(view_region.identity == second_region.identity); + REQUIRE(view_region.offset == second_region.offset + 4 * sizeof(uint32_t)); + REQUIRE(view_region.length == 8 * sizeof(uint32_t)); + REQUIRE(view_region.ne[0] == 8); + + ggml::hrx::ValueBufferBinding binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(second, binding)); + REQUIRE(binding.buffer == nullptr); + REQUIRE(binding.host_data == source.data()); + REQUIRE(binding.offset == second_region.offset); + REQUIRE(binding.weight); + + std::array readback = {}; + ggml_backend_tensor_get(second, readback.data(), 0, sizeof(readback)); + REQUIRE(readback[0] == source[32]); + const uint64_t old_generation = second_region.generation; + const uint32_t replacement = 999; + ggml_backend_tensor_set(second, &replacement, 0, sizeof(replacement)); + REQUIRE(source[32] == replacement); + REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); + REQUIRE(second_region.generation != old_generation); + + const uint32_t async_replacement = 1234; + ggml_backend_tensor_set_async(backend, second, &async_replacement, 0, sizeof(async_replacement)); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get_async(backend, second, readback.data(), 0, sizeof(readback)); + ggml_backend_synchronize(backend); + REQUIRE(readback[0] == async_replacement); + REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); + REQUIRE(second_region.generation != old_generation); + + ggml_backend_buffer_t owned = ggml_backend_alloc_buffer(backend, sizeof(readback)); + REQUIRE(owned != nullptr); + ggml_tensor * copy = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 32); + REQUIRE(ggml_backend_tensor_alloc(owned, copy, ggml_backend_buffer_get_base(owned)) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_copy_async(backend, backend, second, copy); + ggml_backend_synchronize(backend); + ggml_backend_tensor_copy_async(backend, backend, copy, first); + ggml_backend_synchronize(backend); + REQUIRE(source[0] == async_replacement); + + ggml_backend_buffer_free(owned); + ggml_backend_buffer_free(imported); + REQUIRE(source[0] == async_replacement); + ggml_free(ggml); + + auto * hrx_device = backend_context(backend)->device; + const bool previous_direct = hrx_device->use_direct_host_bindings; + hrx_device->use_direct_host_bindings = true; + ggml_backend_buffer_t coherent = ggml_backend_buft_alloc_buffer(ggml_backend_dev_host_buffer_type(device), 1024); + hrx_device->use_direct_host_bindings = previous_direct; + REQUIRE(coherent != nullptr); + auto * coherent_context = ggml_backend_hrx_buffer_context_from_buffer(coherent); + uint8_t * subregion = static_cast(ggml_backend_buffer_get_base(coherent)) + 256; + ggml_backend_buffer_t imported_coherent = ggml_backend_dev_buffer_from_host_ptr(device, subregion, 256, 256); + REQUIRE(imported_coherent != nullptr); + REQUIRE(!ggml_backend_buffer_is_host(imported_coherent)); + auto * imported_context = ggml_backend_hrx_buffer_context_from_buffer(imported_coherent); + REQUIRE(imported_context->imported); + REQUIRE(imported_context->direct_host_binding); + REQUIRE(imported_context->buffer == coherent_context->buffer); + REQUIRE(imported_context->physical_offset == 256); + ggml_init_params subparams = params; + subparams.mem_size = 128 * 1024; + ggml_context * subctx = ggml_init(subparams); + REQUIRE(subctx != nullptr); + ggml_tensor * subtensor = ggml_new_tensor_1d(subctx, GGML_TYPE_I32, 32); + REQUIRE(ggml_backend_tensor_alloc(imported_coherent, subtensor, subregion) == GGML_STATUS_SUCCESS); + ggml::hrx::ValueBufferBinding direct_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(subtensor, direct_binding)); + REQUIRE(direct_binding.buffer == coherent_context->buffer); + REQUIRE(direct_binding.host_data == nullptr); + REQUIRE(direct_binding.offset == 256); + const uint32_t queued_value = 0x456789ab; + require_hrx_status(hrx_stream_fill_buffer(backend_context(backend)->stream, coherent_context->buffer, + 256, 32 * sizeof(uint32_t), &queued_value, sizeof(queued_value))); + std::array ordered_readback = {}; + ggml_backend_tensor_get_async(backend, subtensor, ordered_readback.data(), 0, sizeof(ordered_readback)); + for (uint32_t value : ordered_readback) { + REQUIRE(value == queued_value); + } + const uint32_t ordered_write = 0x13579bdf; + require_hrx_status(hrx_stream_fill_buffer(backend_context(backend)->stream, coherent_context->buffer, + 256, 32 * sizeof(uint32_t), &queued_value, sizeof(queued_value))); + ggml_backend_tensor_set_async(backend, subtensor, &ordered_write, 0, sizeof(ordered_write)); + ggml_backend_synchronize(backend); + REQUIRE(*reinterpret_cast(subregion) == ordered_write); + + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(cpu != nullptr); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_buffer(cpu, sizeof(ordered_readback)); + REQUIRE(cpu_buffer != nullptr); + ggml_context * cpu_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + ggml_tensor * cpu_tensor = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 32); + REQUIRE(ggml_backend_tensor_alloc(cpu_buffer, cpu_tensor, ggml_backend_buffer_get_base(cpu_buffer)) == GGML_STATUS_SUCCESS); + std::array cpu_values = {}; + cpu_values.fill(0x2468ace0); + ggml_backend_tensor_set(cpu_tensor, cpu_values.data(), 0, sizeof(cpu_values)); + require_hrx_status(hrx_stream_fill_buffer(backend_context(backend)->stream, coherent_context->buffer, + 256, sizeof(cpu_values), &queued_value, sizeof(queued_value))); + ggml_backend_tensor_copy_async(cpu, backend, cpu_tensor, subtensor); + ggml_backend_synchronize(backend); + REQUIRE(*reinterpret_cast(subregion) == cpu_values[0]); + + ggml_tensor * coherent_output = ggml_new_tensor_1d(subctx, GGML_TYPE_F32, 32); + REQUIRE(ggml_backend_tensor_alloc(imported_coherent, coherent_output, subregion) == GGML_STATUS_SUCCESS); + ggml_tensor * coherent_input = ggml_new_tensor_1d(subctx, GGML_TYPE_F32, 32); + ggml_backend_buffer_t coherent_input_buffer = ggml_backend_alloc_buffer(backend, 32 * sizeof(float)); + REQUIRE(coherent_input_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(coherent_input_buffer, coherent_input, + ggml_backend_buffer_get_base(coherent_input_buffer)) == GGML_STATUS_SUCCESS); + std::array coherent_values = {}; + coherent_values.fill(3.25f); + ggml_backend_tensor_set(coherent_input, coherent_values.data(), 0, sizeof(coherent_values)); + ggml_tensor * coherent_copy = ggml_cpy(subctx, coherent_input, coherent_output); + ggml_cgraph * coherent_graph = ggml_new_graph(subctx); + REQUIRE(coherent_graph != nullptr); + ggml_build_forward_expand(coherent_graph, coherent_copy); + coherent_graph->uid = 950005; + ggml_backend_hrx_imported_region coherent_region = {}; + REQUIRE(ggml_backend_hrx_get_imported_region(coherent_output, coherent_region)); + const uint64_t coherent_generation = coherent_region.generation; + REQUIRE(ggml_backend_graph_compute(backend, coherent_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + REQUIRE(ggml_backend_hrx_get_imported_region(coherent_output, coherent_region)); + REQUIRE(coherent_region.generation != coherent_generation); + REQUIRE(*reinterpret_cast(subregion) == coherent_values[0]); + ggml_backend_buffer_free(coherent_input_buffer); + ggml_backend_buffer_free(cpu_buffer); + ggml_free(cpu_ctx); + ggml_backend_free(cpu); + ggml_backend_buffer_free(imported_coherent); + ggml_backend_buffer_free(coherent); + ggml_free(subctx); + +#if defined(__linux__) + // Imported file mappings can be read-only. Import, tensor initialization and buffer release must only borrow it. + constexpr size_t page_size = 4096; + void * protected_page = mmap(nullptr, page_size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + REQUIRE(protected_page != MAP_FAILED); + auto * protected_bytes = static_cast(protected_page); + protected_bytes[0] = 0x5a; + REQUIRE(mprotect(protected_page, page_size, PROT_READ) == 0); + ggml_backend_buffer_t read_only_import = ggml_backend_dev_buffer_from_host_ptr( + device, protected_page, page_size, page_size); + REQUIRE(read_only_import != nullptr); + REQUIRE(!ggml_backend_buffer_is_host(read_only_import)); + ggml_context * read_only_ctx = ggml_init(params); + REQUIRE(read_only_ctx != nullptr); + ggml_tensor * read_only_tensor = ggml_new_tensor_1d(read_only_ctx, GGML_TYPE_I8, 1); + REQUIRE(ggml_backend_tensor_alloc(read_only_import, read_only_tensor, protected_page) == GGML_STATUS_SUCCESS); + uint8_t protected_readback = 0; + ggml_backend_tensor_get(read_only_tensor, &protected_readback, 0, sizeof(protected_readback)); + REQUIRE(protected_readback == 0x5a); + ggml_backend_buffer_free(read_only_import); + REQUIRE(protected_bytes[0] == 0x5a); + ggml_free(read_only_ctx); + REQUIRE(munmap(protected_page, page_size) == 0); +#endif +} + static void run_host_transfer_checks(ggml_backend_hrx_context * context) { ggml::hrx::HostTransferManager transfers; ggml::hrx::HostStagingBuffer staging; @@ -345,12 +589,103 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(transfer_stats.uploads == 3); REQUIRE(transfer_stats.upload_bytes == 96); + ggml::hrx::RecordedCommandGraph recorded; + hrx_buffer_retain(first.lease.buffer()); + recorded.retained_buffers.push_back(first.lease.buffer()); + recorded.resident_host_weights.push_back(first.lease); weights.clear(); weight_stats = weights.stats(); - REQUIRE(weight_stats.hits == 0); - REQUIRE(weight_stats.misses == 0); + REQUIRE(weight_stats.hits == 1); + REQUIRE(weight_stats.misses == 3); + REQUIRE(weight_stats.allocation_count == 3); + REQUIRE(weight_stats.resident_bytes == 96); + first.lease = {}; + second.lease = {}; + slice.lease = {}; + next_generation.lease = {}; + weight_stats = weights.stats(); + REQUIRE(weight_stats.allocation_count == 1); + REQUIRE(weight_stats.resident_bytes == 32); + recorded = {}; + weight_stats = weights.stats(); REQUIRE(weight_stats.allocation_count == 0); REQUIRE(weight_stats.resident_bytes == 0); + + source.layout = ggml::hrx::kNativeWeightLayout; + source.generation = 3; + source.offset = 16; + source.source_type = GGML_TYPE_F32; + source.source_ne = { 8, 1, 1, 1 }; + source.source_nb = { 4, 32, 32, 32 }; + auto region = weights.acquire(context->device->device, context->stream, transfers, source); + auto alias = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(region.valid() && alias.valid()); + REQUIRE(region.lease.buffer() == alias.lease.buffer()); + source.source_type = GGML_TYPE_I32; + auto other_type = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(other_type.valid()); + REQUIRE(other_type.lease.buffer() != region.lease.buffer()); + source.source_type = GGML_TYPE_F32; + source.source_ne = { 4, 2, 1, 1 }; + auto other_shape = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(other_shape.valid()); + REQUIRE(other_shape.lease.buffer() != region.lease.buffer()); + source.source_ne = { 8, 1, 1, 1 }; + source.source_nb = { 4, 64, 64, 64 }; + auto other_stride = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(other_stride.valid()); + REQUIRE(other_stride.lease.buffer() != region.lease.buffer()); + source.source_nb = { 4, 32, 32, 32 }; + source.offset = 32; + auto other_slice = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(other_slice.valid()); + REQUIRE(other_slice.lease.buffer() != region.lease.buffer()); + weight_stats = weights.stats(); + REQUIRE(weight_stats.hits == 2); + REQUIRE(weight_stats.misses == 8); + REQUIRE(weight_stats.allocation_count == 5); + REQUIRE(weight_stats.peak_resident_bytes == 160); +} + +static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { + constexpr size_t weight_bytes = 64 * 210; + std::vector host(weight_bytes, 0x5a); + ggml::hrx::HostTransferManager transfers; + ggml::hrx::HostWeightCache weights; + ggml::hrx::HostWeightSource source; + source.host_data = host.data(); + source.identity = 0x5678; + source.generation = 1; + source.capacity = weight_bytes; + source.length = weight_bytes; + source.materialized_length = weight_bytes; + source.layout = ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + source.source_type = GGML_TYPE_Q6_K; + source.input_size = 256; + source.output_size = 64; + + auto first = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(first.valid()); + std::vector downloaded(weight_bytes); + require_hrx_status(hrx_synchronous_d2h(context->device->device, first.lease.buffer(), 0, + downloaded.data(), downloaded.size())); + for (uint8_t byte : downloaded) { + REQUIRE(byte == 0x5a); + } + + first.lease = {}; + weights.evict_except({}, {}); + REQUIRE(weights.stats().allocation_count == 0); + std::fill(host.begin(), host.end(), 0xa5); + source.generation = 2; + auto second = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(second.valid()); + require_hrx_status(hrx_synchronous_d2h(context->device->device, second.lease.buffer(), 0, + downloaded.data(), downloaded.size())); + for (uint8_t byte : downloaded) { + REQUIRE(byte == 0xa5); + } + REQUIRE(weights.stats().peak_resident_bytes == weight_bytes); } template @@ -373,6 +708,11 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_host_buffer_checks(backend); }); }); + suite.device_case("imported_buffer", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_imported_buffer_checks(backend); + }); + }); suite.device_case("host_transfer", [] { run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { run_host_transfer_checks(context); @@ -388,6 +728,11 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_host_weight_cache_checks(context); }); }); + suite.device_case("q6_pack_scratch", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_q6_pack_scratch_checks(context); + }); + }); } int main(int argc, char ** argv) { diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index e8cb79bac47b..128194363ab3 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -1,9 +1,13 @@ #include "backend-context.h" +#include "backend-buffer-binding.h" #include "dispatch/dispatch-scheduler.h" #include "dispatch_registration/common/dispatch-activation-publication.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml-hrx.h" +#include "hrx-interop-utils.h" +#include "hrx_runtime.h" +#include "ggml-impl.h" #include "ggml-quants.h" #include "ggml.h" #include "graph/graph.h" @@ -14,15 +18,19 @@ #include #include #include +#include #include +#include #include #include #include #include #include #include +#include #include #include +#include #include #include @@ -397,6 +405,50 @@ static std::vector scheduled_kernel_sequence(ggml_cgraph * graph) { return names; } +static void require_native_weight_bindings(ggml_cgraph * graph, const std::vector & weights) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + for (ggml_tensor * weight : weights) { + REQUIRE(!ggml_backend_hrx_weight_has_repackable_source(weight)); + const ggml::hrx::Value * value = imported.graph.values().find_tensor(weight); + REQUIRE(value != nullptr); + size_t binding_count = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + ++binding_count; + } + } + } + REQUIRE(binding_count > 0); + } +} + +static void require_packed_weight_bindings(ggml_cgraph * graph, const std::vector & weights) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + for (ggml_tensor * weight : weights) { + REQUIRE(ggml_backend_hrx_weight_has_repackable_source(weight)); + const ggml::hrx::Value * value = imported.graph.values().find_tensor(weight); + REQUIRE(value != nullptr); + size_t binding_count = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == ggml::hrx::kQ4KPackedK256Row64Layout); + ++binding_count; + } + } + } + REQUIRE(binding_count > 0); + } +} + static ggml::hrx::KernelSpecialization scheduled_kernel_specialization(ggml_cgraph * graph, const char * kernel_name) { ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); @@ -2762,6 +2814,2167 @@ static void run_router_projection_case(int64_t token_count) { ggml_backend_free(backend); } +static void run_imported_native_weight_case() { + static constexpr int64_t kInputSize = 2048; + static constexpr int64_t kOutputSize = 128; + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * backend_state = static_cast(backend->context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kInputSize, kOutputSize); + ggml_tensor * weight_alias = ggml_view_2d(ctx, weight, kInputSize, kOutputSize, weight->nb[1], 0); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kInputSize, 1); + ggml_tensor * output = ggml_mul_mat(ctx, weight_alias, input); + REQUIRE(output != nullptr); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + graph->uid = 950001; + + std::vector weight_data = make_router_weight(kInputSize, kOutputSize); + const std::vector input_data = make_router_input(kInputSize, 1); + const std::vector expected = + router_projection_reference(input_data, weight_data, kInputSize, kOutputSize, 1); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), weight_data.data(), weight_data.size() * sizeof(float), + weight_data.size() * sizeof(float)); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, weight, weight_data.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t transient = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(transient != nullptr); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + + ggml_backend_hrx_imported_region region = {}; + REQUIRE(ggml_backend_hrx_get_imported_region(weight, region)); + REQUIRE(region.type == GGML_TYPE_F32); + REQUIRE(region.ne[0] == kInputSize); + REQUIRE(region.ne[1] == kOutputSize); + REQUIRE(region.offset == 0); + ggml_backend_hrx_imported_region alias_region = {}; + REQUIRE(ggml_backend_hrx_get_imported_region(weight_alias, alias_region)); + REQUIRE(alias_region.identity == region.identity); + REQUIRE(alias_region.offset == region.offset); + REQUIRE(alias_region.length == region.length); + const auto before = backend_state->host_weights.stats(); + for (int run = 0; run < 2; ++run) { + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + require_close(actual, expected, 1.0e-2f); + } + const auto after = backend_state->host_weights.stats(); + REQUIRE(after.misses == before.misses + 1); + REQUIRE(after.allocation_count == before.allocation_count + 1); + ggml_backend_hrx_cache_stats cache = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache)); + REQUIRE(cache.prepared_program_hits >= 1); + std::vector readback(weight_data.size()); + ggml_backend_tensor_get(weight, readback.data(), 0, readback.size() * sizeof(float)); + REQUIRE(readback == weight_data); + + ggml_backend_buffer_free(transient); + ggml_backend_buffer_free(imported); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_mutable_host_input_after_writeback_case(bool hrx_host) { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * backend_state = static_cast(backend->context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + constexpr int64_t count = 128; + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, count); + ggml_tensor * written = ggml_scale_inplace(ctx, input, 1.0f); + ggml_tensor * output = ggml_scale(ctx, input, 2.0f); + ggml_cgraph * writer = ggml_new_graph(ctx); + ggml_cgraph * reader = ggml_new_graph(ctx); + REQUIRE(written != nullptr); + REQUIRE(output != nullptr); + REQUIRE(writer != nullptr); + REQUIRE(reader != nullptr); + ggml_build_forward_expand(writer, written); + ggml_build_forward_expand(reader, output); + writer->uid = 950011; + reader->uid = 950012; + + ggml_backend_buffer_type_t host_type = hrx_host ? + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend)) : ggml_backend_cpu_buffer_type(); + ggml_backend_buffer_t host = ggml_backend_buft_alloc_buffer(host_type, ggml_nbytes(input)); + REQUIRE(host != nullptr); + REQUIRE(ggml_backend_tensor_alloc(host, input, ggml_backend_buffer_get_base(host)) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_t device = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(device != nullptr); + + const std::vector previous = make_pattern_f32(count, 21, 0.01f); + const std::vector current = make_pattern_f32(count, 42, 0.03f); + std::vector expected = current; + for (float & value : expected) { + value *= 2.0f; + } + ggml_backend_hrx_cache_stats cold = {}; + for (int run = 0; run < 2; ++run) { + ggml_backend_tensor_set(input, previous.data(), 0, ggml_nbytes(input)); + hrx_timeline_point_t before_writer = {}; + REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(backend_state->stream, &before_writer))); + REQUIRE(ggml_backend_graph_compute_async(backend, writer) == GGML_STATUS_SUCCESS); + REQUIRE(backend_state->graph_replay_state.has_pending_host_writebacks()); + hrx_timeline_point_t after_writer = {}; + REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(backend_state->stream, &after_writer))); + const char * replay_setting = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool replay_enabled = replay_setting == nullptr || replay_setting[0] == '0'; + if (replay_enabled) { + // The recorded launch advances once; the queued host-download tail must advance again. + REQUIRE(after_writer.value >= before_writer.value + 2); + } + REQUIRE(backend_state->weight_submissions.poll_completed().success()); + REQUIRE(backend_state->graph_replay_state.has_pending_host_writebacks()); + // Finish reading registered host memory without draining the backend's deferred writebacks. + REQUIRE(!ggml::hrx::take_status(hrx_stream_synchronize(backend_state->stream))); + REQUIRE(backend_state->graph_replay_state.has_pending_host_writebacks()); + // Qwen's metadata initialization writes staged inputs; the next microbatch supplies new host values. + ggml_backend_tensor_set(input, current.data(), 0, ggml_nbytes(input)); + REQUIRE(ggml_backend_graph_compute_async(backend, reader) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, output), expected, 1.0e-6f); + if (run == 0) { + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cold)); + } else { + ggml_backend_hrx_cache_stats warm = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warm)); + REQUIRE(warm.graph_program_builds == cold.graph_program_builds); + REQUIRE(warm.prepared_program_builds == cold.prepared_program_builds); + REQUIRE(warm.prepared_program_hits >= cold.prepared_program_hits + 2); + } + } + + ggml_backend_buffer_free(device); + ggml_backend_buffer_free(host); + ggml_free(ctx); + ggml_backend_free(backend); +} + +enum class PackedWeightSourceKind { Imported, CpuHost, HrxHost, CoherentSubregion }; + +static void run_weight_demand_comparison_case() { + ggml::hrx::WeightSubmissionTracker tracker; + ggml::hrx::HostWeightSource native; + native.identity = 17; + native.generation = 3; + native.capacity = 65536; + native.offset = 4096; + native.length = 32768; + native.region_length = 32768; + native.region_type = GGML_TYPE_F16; + native.materialized_length = 32768; + native.source_type = GGML_TYPE_F16; + native.input_size = 256; + native.output_size = 64; + native.source_ne = { 256, 64, 1, 1 }; + native.source_nb = { 2, 512, 32768, 32768 }; + ggml::hrx::HostWeightSource packed = native; + packed.layout = "packed-test-layout"; + packed.materialized_length = 49152; + ggml::hrx::HostWeightSource other = native; + other.offset = 36864; + + REQUIRE(tracker.demand_matches({})); + REQUIRE(!tracker.demand_matches({ native })); + REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native, other }, true, true).success()); + REQUIRE(tracker.demand_matches({})); + REQUIRE(tracker.demand_matches({ native })); + int aliases[2] = {}; + ggml::hrx::HostWeightSource alias = native; + alias.host_data = &aliases[1]; + REQUIRE(tracker.demand_matches({ native, alias, native })); + REQUIRE(tracker.demand_matches({ other, other })); + REQUIRE(!tracker.demand_matches({ native, packed })); + REQUIRE(!tracker.demand_matches({ packed, native })); + REQUIRE(!tracker.demand_matches({ packed })); + + auto expect_mismatch = [&](const ggml::hrx::HostWeightSource & changed) { + REQUIRE(!tracker.demand_matches({ changed })); + }; + ggml::hrx::HostWeightSource changed = native; + changed.generation++; + expect_mismatch(changed); + changed = native; + changed.capacity++; + expect_mismatch(changed); + changed = native; + changed.length++; + expect_mismatch(changed); + changed = native; + changed.materialized_length++; + expect_mismatch(changed); + changed = native; + changed.source_type = GGML_TYPE_F32; + expect_mismatch(changed); + changed = native; + changed.input_size++; + expect_mismatch(changed); + changed = native; + changed.output_size++; + expect_mismatch(changed); + changed = native; + changed.identity++; + expect_mismatch(changed); + changed = native; + changed.offset++; + expect_mismatch(changed); + changed = native; + changed.region_length++; + expect_mismatch(changed); + changed = native; + changed.region_type = GGML_TYPE_F32; + expect_mismatch(changed); + changed = native; + changed.source_ne[0]++; + expect_mismatch(changed); + changed = native; + changed.source_nb[0]++; + expect_mismatch(changed); + + REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native, packed, other }, true, true).success()); + REQUIRE(tracker.demand_matches({ packed, native, packed, native })); + REQUIRE(tracker.demand_matches({ other, native, packed })); + REQUIRE(!tracker.demand_matches({ native })); + ggml::hrx::HostWeightSource third = packed; + third.layout = "another-packed-layout"; + REQUIRE(!tracker.demand_matches({ native, packed, third })); + ggml::hrx::HostWeightSource changed_other = other; + changed_other.generation++; + REQUIRE(!tracker.demand_matches({ changed_other, native, packed })); + REQUIRE(!tracker.demand_matches({ native, packed, changed_other })); +} + +static void run_weight_reclamation_sequence_case() { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(cpu != nullptr && backend != nullptr); + auto * state = static_cast(backend->context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_second_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_second_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_third_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_third_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * cpu_second_rows = ggml_get_rows(cpu_ctx, cpu_second_weight, cpu_ids); + ggml_tensor * cpu_third_rows = ggml_get_rows(cpu_ctx, cpu_third_weight, cpu_ids); + ggml_tensor * cpu_matmul = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * cpu_third_matmul = ggml_mul_mat(cpu_ctx, cpu_third_weight, cpu_input); + ggml_tensor * cpu_no_weight_output = ggml_add(cpu_ctx, cpu_input, cpu_input); + ggml_cgraph * cpu_native = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_packed = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_third_packed = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_second_native = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_no_weight = ggml_new_graph(cpu_ctx); + ggml_build_forward_expand(cpu_native, cpu_rows); + ggml_build_forward_expand(cpu_packed, cpu_matmul); + ggml_build_forward_expand(cpu_third_packed, cpu_third_matmul); + ggml_build_forward_expand(cpu_second_native, cpu_second_rows); + ggml_build_forward_expand(cpu_no_weight, cpu_no_weight_output); + ggml_cgraph * cpu_two_weights = ggml_new_graph(cpu_ctx); + ggml_build_forward_expand(cpu_two_weights, cpu_rows); + ggml_build_forward_expand(cpu_two_weights, cpu_third_rows); + + ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_matmul = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_third_matmul = ggml_mul_mat(hrx_ctx, hrx_third_weight, hrx_input); + ggml_tensor * hrx_rows_structural = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_matmul_structural = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_second_rows = ggml_get_rows(hrx_ctx, hrx_second_weight, hrx_ids); + ggml_tensor * hrx_third_rows = ggml_get_rows(hrx_ctx, hrx_third_weight, hrx_ids); + ggml_tensor * hrx_two_first = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_two_second = ggml_get_rows(hrx_ctx, hrx_third_weight, hrx_ids); + ggml_tensor * hrx_no_weight_output = ggml_add(hrx_ctx, hrx_input, hrx_input); + ggml_cgraph * hrx_native = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_packed = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_third_packed = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_native_structural = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_packed_structural = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_second_native = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_empty = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_no_weight = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_two_weights = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_two_weights_structural = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(hrx_native, hrx_rows); + ggml_build_forward_expand(hrx_packed, hrx_matmul); + ggml_build_forward_expand(hrx_third_packed, hrx_third_matmul); + ggml_build_forward_expand(hrx_native_structural, hrx_rows_structural); + ggml_build_forward_expand(hrx_packed_structural, hrx_matmul_structural); + ggml_build_forward_expand(hrx_second_native, hrx_second_rows); + ggml_build_forward_expand(hrx_no_weight, hrx_no_weight_output); + ggml_build_forward_expand(hrx_two_weights, hrx_rows); + ggml_build_forward_expand(hrx_two_weights, hrx_third_rows); + ggml_build_forward_expand(hrx_two_weights_structural, hrx_two_first); + ggml_build_forward_expand(hrx_two_weights_structural, hrx_two_second); + hrx_native->uid = 951061; + hrx_packed->uid = 951062; + hrx_third_packed->uid = 951069; + hrx_native_structural->uid = 951063; + hrx_packed_structural->uid = 951064; + hrx_second_native->uid = 951067; + hrx_no_weight->uid = 951068; + hrx_two_weights->uid = 951065; + hrx_two_weights_structural->uid = 951066; + + std::vector weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 11); + std::vector second_weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 23); + std::vector source_storage(weight.size() + second_weight.size()); + std::copy(weight.begin(), weight.end(), source_storage.begin()); + std::copy(second_weight.begin(), second_weight.end(), source_storage.begin() + weight.size()); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), source_storage.data(), source_storage.size(), source_storage.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, source_storage.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_second_weight, source_storage.data() + weight.size()) == + GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + std::vector third_weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 47); + ggml_backend_buffer_t third_imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), third_weight.data(), third_weight.size(), third_weight.size()); + REQUIRE(third_imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(third_imported, hrx_third_weight, third_weight.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(third_imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); + REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + ggml_backend_tensor_set(cpu_second_weight, second_weight.data(), 0, second_weight.size()); + ggml_backend_tensor_set(cpu_third_weight, third_weight.data(), 0, third_weight.size()); + const std::array ids = { 0, 1, 63 }; + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); + const std::vector input = make_pattern_f32(input_size * 3, 17, 0.01f); + set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_native) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_packed) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_third_packed) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_second_native) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_no_weight) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_two_weights) == GGML_STATUS_SUCCESS); + const std::vector expected_native = get_f32_tensor(cpu, cpu_rows); + const std::vector expected_packed = get_f32_tensor(cpu, cpu_matmul); + const std::vector expected_third_packed = get_f32_tensor(cpu, cpu_third_matmul); + const std::vector expected_second = get_f32_tensor(cpu, cpu_second_rows); + const std::vector expected_third = get_f32_tensor(cpu, cpu_third_rows); + const std::vector expected_no_weight = get_f32_tensor(cpu, cpu_no_weight_output); + float source_gap = 0.0f; + for (size_t index = 0; index < expected_native.size(); ++index) { + source_gap = std::max(source_gap, std::fabs(expected_native[index] - expected_second[index])); + } + REQUIRE(source_gap > 0.1f); + REQUIRE(std::any_of(expected_packed.begin(), expected_packed.end(), + [](float value) { return std::fabs(value) > 0.1f; })); + + auto execute = [&](ggml_cgraph * graph, ggml_tensor * output, const std::vector & expected, + size_t live_allocations = 1) { + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, output), expected, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.weight_live_allocations == live_allocations); + return stats; + }; + const auto native_first = execute(hrx_native, hrx_rows, expected_native); + REQUIRE(native_first.weight_materializations == 1); + const auto packed_first = execute(hrx_packed, hrx_matmul, expected_packed); + REQUIRE(packed_first.weight_materializations == 2); + REQUIRE(packed_first.weight_peak_bytes == + std::max(native_first.weight_live_bytes, packed_first.weight_live_bytes)); + const auto before_native_rebuild = state->graph_programs.stats(); + const auto native_second = execute(hrx_native, hrx_rows, expected_native); + REQUIRE(native_second.weight_materializations == 3); + const auto rebuilt_native_demand = state->graph_programs.stats(); + REQUIRE(rebuilt_native_demand.weight_demand_builds > before_native_rebuild.weight_demand_builds); + const auto native_unchanged = execute(hrx_native, hrx_rows, expected_native); + REQUIRE(native_unchanged.weight_materializations == native_second.weight_materializations); + REQUIRE(native_unchanged.weight_uploaded_bytes == native_second.weight_uploaded_bytes); + REQUIRE(native_unchanged.prepared_program_hits > native_second.prepared_program_hits); + REQUIRE(state->graph_programs.stats().weight_demand_hits > rebuilt_native_demand.weight_demand_hits); + hrx_timeline_point_t before_replay = {}; + REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &before_replay))); + const auto recorded_native = ggml::hrx::GraphExecutor(*state).execute(*hrx_native); + REQUIRE(recorded_native.success()); + hrx_timeline_point_t after_replay = {}; + REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &after_replay))); + const char * replay_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool replay_disabled = replay_flag != nullptr && std::string(replay_flag) == "1"; + REQUIRE(recorded_native.graph_replay_event == (replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : + ggml::hrx::HrxGraphReplayEvent::Hit)); + if (!replay_disabled) { + REQUIRE(after_replay.semaphore == before_replay.semaphore); + REQUIRE(after_replay.value == before_replay.value + 1); + } + ggml::hrx::WeightSubmissionTracker isolated_tracker; + auto isolated_resources = std::make_shared(); + REQUIRE(isolated_tracker.record(state->device->device, state->stream, isolated_resources, {}, true, false).success()); + hrx_timeline_point_t after_tracker = {}; + REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &after_tracker))); + REQUIRE(after_tracker.semaphore == after_replay.semaphore); + REQUIRE(after_tracker.value == after_replay.value); + ggml_backend_synchronize(backend); + + const auto structural_before = state->prepared_programs.stats(); + const auto structural_native = execute(hrx_native_structural, hrx_rows_structural, expected_native); + REQUIRE(structural_native.weight_materializations == native_second.weight_materializations); + const auto structural_packed = execute(hrx_packed_structural, hrx_matmul_structural, expected_packed); + REQUIRE(structural_packed.weight_materializations == native_second.weight_materializations + 1); + const auto structural_return = execute(hrx_native_structural, hrx_rows_structural, expected_native); + REQUIRE(structural_return.weight_materializations == native_second.weight_materializations + 2); + REQUIRE(state->prepared_programs.stats().builds > structural_before.builds); + const auto distinct_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); + REQUIRE(distinct_second.weight_materializations == structural_return.weight_materializations + 1); + REQUIRE(ggml_backend_graph_compute(backend, hrx_empty) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + REQUIRE(hrx_no_weight->n_nodes > 0); + const auto no_weight = execute(hrx_no_weight, hrx_no_weight_output, expected_no_weight, 2); + REQUIRE(no_weight.weight_materializations == distinct_second.weight_materializations); + REQUIRE(no_weight.weight_uploaded_bytes == distinct_second.weight_uploaded_bytes); + REQUIRE(no_weight.weight_live_bytes == distinct_second.weight_live_bytes); + const auto reacquired_native = execute(hrx_native, hrx_rows, expected_native, 2); + REQUIRE(reacquired_native.weight_materializations == no_weight.weight_materializations); + REQUIRE(reacquired_native.weight_uploaded_bytes == no_weight.weight_uploaded_bytes); + const auto repeated_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); + REQUIRE(repeated_second.weight_materializations == reacquired_native.weight_materializations); + REQUIRE(repeated_second.weight_uploaded_bytes == reacquired_native.weight_uploaded_bytes); + const auto rerecorded_native = ggml::hrx::GraphExecutor(*state).execute(*hrx_native); + REQUIRE(rerecorded_native.success()); + REQUIRE(rerecorded_native.graph_replay_event == (replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : + ggml::hrx::HrxGraphReplayEvent::Hit)); + ggml_backend_synchronize(backend); + const auto packed_with_second = execute(hrx_packed, hrx_matmul, expected_packed, 2); + const auto second_between_packed = execute(hrx_second_native, hrx_second_rows, expected_second, 2); + const auto packed_reused = execute(hrx_packed, hrx_matmul, expected_packed, 2); + REQUIRE(second_between_packed.weight_materializations == packed_with_second.weight_materializations); + REQUIRE(packed_reused.weight_materializations == packed_with_second.weight_materializations); + REQUIRE(packed_reused.weight_uploaded_bytes == packed_with_second.weight_uploaded_bytes); + auto verify_two_weights = [&](ggml_cgraph * graph, ggml_tensor * first, ggml_tensor * second) { + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, first), expected_native, 2.0e-2f, 3.0e-2f); + require_close(get_f32_tensor(backend, second), expected_third, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.weight_live_allocations == 3); + }; + auto fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second, + size_t warm_allocations) { + const auto before_cold_failure = execute(hrx_packed, hrx_matmul, expected_packed, warm_allocations); + state->host_weights.fail_cold_acquire_after_for_test(1); + const auto before_failure = state->graph_programs.stats(); + const auto retained_before_failure = state->weight_submissions.current_demand(); + REQUIRE(ggml_backend_graph_compute(backend, two_weights) == GGML_STATUS_FAILED); + REQUIRE(state->graph_programs.stats().weight_demand_builds == before_failure.weight_demand_builds); + REQUIRE(state->weight_submissions.demand_matches(retained_before_failure)); + ggml_backend_synchronize(backend); + ggml_backend_hrx_cache_stats failed = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &failed)); + REQUIRE(failed.weight_live_allocations == 1); + REQUIRE(failed.weight_materializations == before_cold_failure.weight_materializations + 1); + const auto preserved_second = execute(hrx_second_native, hrx_second_rows, expected_second, 1); + REQUIRE(preserved_second.weight_materializations == failed.weight_materializations); + execute(hrx_packed, hrx_matmul, expected_packed, 2); + verify_two_weights(two_weights, first, second); + }; + fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows, 2); + // A layout transition makes the next native acquisition cold without changing the immutable source. + execute(hrx_third_packed, hrx_third_matmul, expected_third_packed, 3); + const auto structural_failure_before = state->prepared_programs.stats(); + fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second, 3); + REQUIRE(state->prepared_programs.stats().builds > structural_failure_before.builds); + verify_two_weights(hrx_two_weights, hrx_rows, hrx_third_rows); + + auto completion_fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second) { + execute(hrx_packed, hrx_matmul, expected_packed, 3); + state->weight_submissions.fail_completion_record_for_test(); + const auto before_failure = state->graph_programs.stats(); + const auto retained_before_failure = state->weight_submissions.current_demand(); + REQUIRE(ggml_backend_graph_compute(backend, two_weights) == GGML_STATUS_FAILED); + REQUIRE(state->graph_programs.stats().weight_demand_builds == before_failure.weight_demand_builds); + REQUIRE(state->weight_submissions.demand_matches(retained_before_failure)); + // The fallback synchronization completed the submitted writes despite the error return. + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, first), expected_native, 2.0e-2f, 3.0e-2f); + require_close(get_f32_tensor(backend, second), expected_third, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats failed = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &failed)); + REQUIRE(failed.weight_live_allocations == 2); + const auto preserved_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); + REQUIRE(preserved_second.weight_materializations == failed.weight_materializations); + execute(hrx_packed, hrx_matmul, expected_packed, 3); + verify_two_weights(two_weights, first, second); + }; + completion_fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows); + completion_fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second); + + ggml_backend_buffer_free(hrx_tensors); + ggml_backend_buffer_free(cpu_tensors); + ggml_backend_buffer_free(imported); + ggml_backend_buffer_free(third_imported); + ggml_free(hrx_ctx); + ggml_free(cpu_ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); +} + +static void run_weight_reclamation_overlap_case() { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool had_flag = old_flag != nullptr; + const std::string old_flag_value = had_flag ? old_flag : ""; + REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "1", 1) == 0); + + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(cpu != nullptr && backend != nullptr); + auto * state = static_cast(backend->context); + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + ggml_tensor * cpu_w0 = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_w1 = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_disjoint_w = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_w0 = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_w1 = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_disjoint_w = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_w0, cpu_ids); + ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_w0, hrx_ids); + ggml_tensor * cpu_disjoint_rows = ggml_get_rows(cpu_ctx, cpu_disjoint_w, cpu_ids); + ggml_tensor * hrx_disjoint_rows = ggml_get_rows(hrx_ctx, hrx_disjoint_w, hrx_ids); + ggml_tensor * cpu_mixed = ggml_mul_mat(cpu_ctx, cpu_w1, cpu_rows); + ggml_tensor * hrx_mixed = ggml_mul_mat(hrx_ctx, hrx_w1, hrx_rows); + ggml_tensor * cpu_packed = ggml_mul_mat(cpu_ctx, cpu_w1, cpu_input); + ggml_tensor * hrx_packed = ggml_mul_mat(hrx_ctx, hrx_w1, hrx_input); + ggml_cgraph * cpu_mixed_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_native_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_packed_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_disjoint_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_mixed_graph = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_native_graph = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_packed_graph = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_disjoint_graph = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(cpu_mixed_graph, cpu_mixed); + ggml_build_forward_expand(cpu_native_graph, cpu_rows); + ggml_build_forward_expand(cpu_packed_graph, cpu_packed); + ggml_build_forward_expand(cpu_disjoint_graph, cpu_disjoint_rows); + ggml_build_forward_expand(hrx_mixed_graph, hrx_mixed); + ggml_build_forward_expand(hrx_native_graph, hrx_rows); + ggml_build_forward_expand(hrx_packed_graph, hrx_packed); + ggml_build_forward_expand(hrx_disjoint_graph, hrx_disjoint_rows); + hrx_mixed_graph->uid = 951071; + hrx_native_graph->uid = 951072; + hrx_packed_graph->uid = 951073; + hrx_disjoint_graph->uid = 951074; + + std::vector weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 31); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), weight.data(), weight.size(), weight.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_w0, weight.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_w1, weight.data()) == GGML_STATUS_SUCCESS); + std::vector disjoint_weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 43); + ggml_backend_buffer_t disjoint_imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), disjoint_weight.data(), disjoint_weight.size(), disjoint_weight.size()); + REQUIRE(disjoint_imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(disjoint_imported, hrx_disjoint_w, disjoint_weight.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_set_usage(disjoint_imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); + REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); + ggml_backend_tensor_set(cpu_w0, weight.data(), 0, weight.size()); + ggml_backend_tensor_set(cpu_w1, weight.data(), 0, weight.size()); + ggml_backend_tensor_set(cpu_disjoint_w, disjoint_weight.data(), 0, disjoint_weight.size()); + const std::array ids = { 0, 1, 63 }; + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); + const std::vector input = make_pattern_f32(input_size * 3, 29, 0.01f); + set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_mixed_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_packed_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_disjoint_graph) == GGML_STATUS_SUCCESS); + const std::vector expected_mixed = get_f32_tensor(cpu, cpu_mixed); + const std::vector expected_native = get_f32_tensor(cpu, cpu_rows); + const std::vector expected_packed = get_f32_tensor(cpu, cpu_packed); + const std::vector expected_disjoint = get_f32_tensor(cpu, cpu_disjoint_rows); + REQUIRE(std::any_of(expected_mixed.begin(), expected_mixed.end(), + [](float value) { return std::fabs(value) > 0.1f; })); + REQUIRE(std::any_of(expected_packed.begin(), expected_packed.end(), + [](float value) { return std::fabs(value) > 0.1f; })); + + auto imported_graph = ggml::hrx::import_ggml_graph(*hrx_mixed_graph); + REQUIRE(imported_graph.valid()); + const auto * w0_value = imported_graph.graph.values().find_tensor(hrx_w0); + const auto * w1_value = imported_graph.graph.values().find_tensor(hrx_w1); + REQUIRE(w0_value != nullptr && w1_value != nullptr); + REQUIRE(w0_value->id != w1_value->id); + + // Preupload the packed layout before closing the gate. Its first coexistence with native + // then occurs across two real submissions, without a cold upload waiting on the gate. + REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + auto packed_lookup = state->graph_programs.get_or_build( + *hrx_packed_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(packed_lookup.valid()); + auto packed_match = packed_lookup.program->match_trusted_graph(*hrx_packed_graph, true); + REQUIRE(packed_match.valid()); + std::vector packed_external; + for (const auto & external : packed_match.external_bindings) { + ggml::hrx::ValueBufferBinding buffer; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = external.tensor->ne[dimension]; + nb[dimension] = external.tensor->nb[dimension]; + } + packed_external.push_back(ggml::hrx::make_command_program_binding( + external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + } + auto packed_bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(packed_external)); + REQUIRE(packed_bindings.valid()); + auto packed_demands = ggml::hrx::collect_command_program_weight_demands( + packed_lookup.program->commands(), packed_bindings); + REQUIRE(packed_demands.valid()); + REQUIRE(packed_demands.weights.size() == 1 && packed_demands.weights.front().requires_materialization); + auto preuploaded = state->host_weights.acquire(state->device->device, state->stream, + state->host_transfers, packed_demands.weights.front().source); + REQUIRE(preuploaded.valid()); + ggml_backend_hrx_cache_stats cross_before = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cross_before)); + REQUIRE(cross_before.weight_duplicate_warnings == 0); + hrx_semaphore_t cross_gate = nullptr; + REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &cross_gate))); + REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { cross_gate, 1 }))); + std::mutex cross_mutex; + std::condition_variable cross_cv; + bool cross_finished = false; + bool cross_timed_out = false; + std::thread cross_watchdog([&] { + std::unique_lock lock(cross_mutex); + if (!cross_cv.wait_for(lock, std::chrono::seconds(5), [&] { return cross_finished; })) { + cross_timed_out = true; + hrx_status_t signal = hrx_semaphore_signal(cross_gate, 1); + hrx_status_ignore(signal); + } + }); + const bool cross_native = ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + const bool cross_packed = cross_native && + ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_SUCCESS; + ggml_backend_hrx_cache_stats cross_queued = {}; + const bool cross_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cross_queued); + bool cross_obsolete_completed = true; + const bool cross_poll_ok = state->weight_submissions.poll_completed(&cross_obsolete_completed).success(); + const bool cross_old_demand_pending = state->weight_submissions.has_obsolete_pending(); + { + std::lock_guard lock(cross_mutex); + cross_finished = true; + } + cross_cv.notify_one(); + const bool cross_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cross_gate, 1)); + cross_watchdog.join(); + REQUIRE(cross_native && cross_packed && cross_stats_ok && cross_poll_ok && cross_release_ok && !cross_timed_out); + REQUIRE(!cross_obsolete_completed && cross_old_demand_pending); + REQUIRE(cross_queued.weight_live_allocations == 2); + REQUIRE(cross_queued.weight_duplicate_warnings == cross_before.weight_duplicate_warnings + 1); + preuploaded.lease = {}; + ggml_backend_synchronize(backend); + REQUIRE(!state->weight_submissions.has_obsolete_pending()); + require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); + hrx_semaphore_release(cross_gate); + + REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats warmed = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); + REQUIRE(warmed.weight_live_allocations == 2); + REQUIRE(warmed.weight_duplicate_warnings >= cross_queued.weight_duplicate_warnings); + REQUIRE(state->prepared_programs.stats().builds == 0); + + // A cold preparation failure must leave an already queued native reader's lease intact. + REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const std::vector poisoned_native(expected_native.size(), 123.0f); + ggml_backend_tensor_set(hrx_rows, poisoned_native.data(), 0, + poisoned_native.size() * sizeof(float)); + hrx_semaphore_t cold_gate = nullptr; + REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &cold_gate))); + REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { cold_gate, 1 }))); + std::mutex cold_mutex; + std::condition_variable cold_cv; + bool cold_finished = false; + bool cold_timed_out = false; + std::thread cold_watchdog([&] { + std::unique_lock lock(cold_mutex); + if (!cold_cv.wait_for(lock, std::chrono::seconds(5), [&] { return cold_finished; })) { + cold_timed_out = true; + hrx_status_t signal = hrx_semaphore_signal(cold_gate, 1); + hrx_status_ignore(signal); + } + }); + const bool native_before_failure = + ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + state->host_weights.fail_cold_acquire_after_for_test(0); + const bool cold_failed = + ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_FAILED; + ggml_backend_hrx_cache_stats cold_pending = {}; + const bool cold_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cold_pending); + { + std::lock_guard lock(cold_mutex); + cold_finished = true; + } + cold_cv.notify_one(); + const bool cold_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cold_gate, 1)); + cold_watchdog.join(); + REQUIRE(native_before_failure && cold_failed && cold_stats_ok && cold_release_ok && !cold_timed_out); + REQUIRE(cold_pending.weight_live_allocations == 1); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); + hrx_semaphore_release(cold_gate); + REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); + REQUIRE(warmed.weight_live_allocations == 2); + + hrx_semaphore_t gate = nullptr; + REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); + REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); + std::mutex gate_mutex; + std::condition_variable gate_cv; + bool finished = false; + bool timed_out = false; + std::thread watchdog([&] { + std::unique_lock lock(gate_mutex); + const bool released = gate_cv.wait_for(lock, std::chrono::seconds(5), [&] { return finished; }); + if (!released) { + timed_out = true; + hrx_status_t signal = hrx_semaphore_signal(gate, 1); + hrx_status_ignore(signal); + } + }); + const bool mixed_queued = ggml_backend_graph_compute_async(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS; + const bool native_queued = mixed_queued && + ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + const bool packed_queued = native_queued && + ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_SUCCESS; + bool stream_complete = true; + const bool stream_query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); + ggml_backend_hrx_cache_stats queued = {}; + const bool stats_ok = ggml_backend_hrx_get_cache_stats(backend, &queued); + { + std::lock_guard lock(gate_mutex); + finished = true; + } + gate_cv.notify_one(); + const bool signal_ok = !ggml::hrx::take_status(hrx_semaphore_signal(gate, 1)); + watchdog.join(); + REQUIRE(mixed_queued && native_queued && packed_queued); + REQUIRE(stream_query_ok && !stream_complete); + REQUIRE(stats_ok); + REQUIRE(queued.weight_live_allocations == 2); + REQUIRE(queued.weight_peak_bytes == warmed.weight_peak_bytes); + REQUIRE(queued.weight_duplicate_warnings == warmed.weight_duplicate_warnings); + REQUIRE(signal_ok); + REQUIRE(!timed_out); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); + require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats retired = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &retired)); + REQUIRE(retired.weight_live_allocations == 1); + REQUIRE(retired.weight_live_bytes < queued.weight_live_bytes); + REQUIRE(retired.weight_duplicate_warnings == queued.weight_duplicate_warnings); + REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats reacquired = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &reacquired)); + REQUIRE(reacquired.weight_materializations == retired.weight_materializations + 1); + REQUIRE(reacquired.weight_live_allocations == 1); + + REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + REQUIRE(ggml_backend_graph_compute(backend, hrx_disjoint_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_disjoint_rows), expected_disjoint, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats mixed_and_disjoint = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &mixed_and_disjoint)); + REQUIRE(mixed_and_disjoint.weight_live_allocations == 3); + REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats mixed_reused = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &mixed_reused)); + REQUIRE(mixed_reused.weight_materializations == mixed_and_disjoint.weight_materializations); + REQUIRE(mixed_reused.weight_live_allocations == 3); + REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_hrx_cache_stats native_narrowed = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &native_narrowed)); + REQUIRE(native_narrowed.weight_live_allocations == 2); + REQUIRE(ggml_backend_graph_compute(backend, hrx_disjoint_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_disjoint_rows), expected_disjoint, 2.0e-2f, 3.0e-2f); + ggml_backend_hrx_cache_stats disjoint_reused = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &disjoint_reused)); + REQUIRE(disjoint_reused.weight_materializations == native_narrowed.weight_materializations); + REQUIRE(disjoint_reused.weight_live_allocations == 2); + + hrx_semaphore_release(gate); + ggml_backend_buffer_free(hrx_tensors); + ggml_backend_buffer_free(cpu_tensors); + ggml_backend_buffer_free(imported); + ggml_backend_buffer_free(disjoint_imported); + ggml_free(hrx_ctx); + ggml_free(cpu_ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); + restore_environment_value("GGML_HRX_DISABLE_GRAPH_REPLAY", had_flag, old_flag_value); +} + +static void run_failed_submission_weight_lifetime_case() { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool had_flag = old_flag != nullptr; + const std::string old_flag_value = had_flag ? old_flag : ""; + REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "1", 1) == 0); + + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(cpu != nullptr && backend != nullptr); + auto * state = static_cast(backend->context); + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_native = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_native_graph = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + ggml_build_forward_expand(hrx_native_graph, hrx_native); + hrx_graph->uid = 951081; + hrx_native_graph->uid = 951082; + + std::vector weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 59); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), weight.data(), weight.size(), weight.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, weight.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); + REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + const std::array ids = { 0, 1, 63 }; + ggml_backend_tensor_set(hrx_ids, ids.data(), 0, sizeof(ids)); + const std::vector input = make_pattern_f32(input_size * 3, 41, 0.01f); + set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + const std::vector expected = get_f32_tensor(cpu, cpu_output); + REQUIRE(std::any_of(expected.begin(), expected.end(), [](float value) { return std::fabs(value) > 0.1f; })); + const std::vector poisoned(expected.size(), 123.0f); + ggml_backend_tensor_set(hrx_output, poisoned.data(), 0, poisoned.size() * sizeof(float)); + REQUIRE(std::any_of(expected.begin(), expected.end(), + [](float value) { return std::fabs(value - 123.0f) > 0.1f; })); + + auto bind = [&](const ggml::hrx::GraphProgram & program, ggml_cgraph & graph) { + auto match = program.match_trusted_graph(graph, true); + REQUIRE(match.valid()); + std::vector external_bindings; + for (const auto & external : match.external_bindings) { + ggml::hrx::ValueBufferBinding buffer; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = external.tensor->ne[dimension]; + nb[dimension] = external.tensor->nb[dimension]; + } + external_bindings.push_back(ggml::hrx::make_command_program_binding( + external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); + REQUIRE(bindings.valid()); + return bindings; + }; + auto packed_lookup = state->graph_programs.get_or_build( + *hrx_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + auto native_lookup = state->graph_programs.get_or_build( + *hrx_native_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(packed_lookup.valid() && native_lookup.valid()); + auto packed_bindings = bind(*packed_lookup.program, *hrx_graph); + auto native_bindings = bind(*native_lookup.program, *hrx_native_graph); + auto packed_demands = ggml::hrx::collect_command_program_weight_demands( + packed_lookup.program->commands(), packed_bindings); + auto native_demands = ggml::hrx::collect_command_program_weight_demands( + native_lookup.program->commands(), native_bindings); + REQUIRE(packed_demands.valid() && native_demands.valid()); + REQUIRE(packed_demands.weights.size() == 1 && packed_demands.weights.front().requires_materialization); + REQUIRE(native_demands.weights.size() == 1); + const auto packed_source = packed_demands.weights.front().source; + const auto native_source = native_demands.weights.front().source; + REQUIRE(native_source.layout == ggml::hrx::kNativeWeightLayout); + REQUIRE(packed_source.layout != native_source.layout); + REQUIRE(packed_source.materialized_length > 0); + ggml::hrx::WeightSubmissionTracker tracker; + REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native_source }, true, true).success()); + REQUIRE(tracker.demand_matches({ native_source })); + REQUIRE(!tracker.demand_matches({ packed_source })); + + auto preuploaded = state->host_weights.acquire(state->device->device, state->stream, + state->host_transfers, packed_source); + REQUIRE(preuploaded.valid()); + ggml_backend_synchronize(backend); + preuploaded.lease = {}; + hrx_semaphore_t gate = nullptr; + REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); + REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); + std::mutex gate_mutex; + std::condition_variable gate_cv; + bool finished = false; + bool timed_out = false; + std::thread watchdog([&] { + std::unique_lock lock(gate_mutex); + if (!gate_cv.wait_for(lock, std::chrono::seconds(5), [&] { return finished; })) { + timed_out = true; + hrx_status_t signal = hrx_semaphore_signal(gate, 1); + hrx_status_ignore(signal); + } + }); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + state->device->device, state->stream, state->device->architecture.c_str(), + &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, + &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, + }; + auto execution = packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); + const bool command_queued = execution.success && execution.execution_resources != nullptr; + std::weak_ptr sole_resource = execution.execution_resources; + const bool has_execution_resources = execution.execution_resources != nullptr; + const bool record_ok = has_execution_resources && tracker.record(state->device->device, state->stream, + execution.execution_resources, { packed_source }, false, true).success(); + state->graph_programs.invalidate_obsolete_weights(tracker.current_demand()); + execution.execution_resources.reset(); + state->host_weights.evict_except(tracker.current_demand(), {}); + const auto pending_stats = state->host_weights.stats(); + bool completed = true; + const bool poll_ok = tracker.poll_completed(&completed).success(); + bool stream_complete = true; + const bool query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); + const bool demand_unchanged = tracker.demand_matches({ native_source }); + const bool tracker_only_owner = sole_resource.use_count() == 1; + const bool resource_retained = tracker.outstanding_leases().size() == 1; + { + std::lock_guard lock(gate_mutex); + finished = true; + } + gate_cv.notify_one(); + const bool signal_ok = !ggml::hrx::take_status(hrx_semaphore_signal(gate, 1)); + watchdog.join(); + REQUIRE(command_queued && record_ok && poll_ok && query_ok && signal_ok && !timed_out); + REQUIRE(!completed && !stream_complete && demand_unchanged && tracker_only_owner && resource_retained); + REQUIRE(pending_stats.allocation_count == 1); + REQUIRE(pending_stats.resident_bytes == packed_source.materialized_length); + + REQUIRE(!ggml::hrx::take_status(hrx_stream_synchronize(state->stream))); + require_close(get_f32_tensor(backend, hrx_output), expected, 2.0e-2f, 3.0e-2f); + bool retired = false; + REQUIRE(tracker.poll_completed(&retired).success()); + REQUIRE(retired); + state->host_weights.evict_except(tracker.current_demand(), tracker.outstanding_leases()); + const auto reclaimed = state->host_weights.stats(); + REQUIRE(sole_resource.expired()); + REQUIRE(tracker.demand_matches({ native_source })); + REQUIRE(reclaimed.allocation_count == 0); + REQUIRE(reclaimed.resident_bytes == 0); + hrx_semaphore_release(gate); + ggml_backend_buffer_free(hrx_tensors); + ggml_backend_buffer_free(cpu_tensors); + ggml_backend_buffer_free(imported); + ggml_free(hrx_ctx); + ggml_free(cpu_ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); + restore_environment_value("GGML_HRX_DISABLE_GRAPH_REPLAY", had_flag, old_flag_value); +} + +static void run_imported_packed_weight_case(ggml_type weight_type, + PackedWeightSourceKind source_kind = PackedWeightSourceKind::Imported) { + static constexpr int64_t kInputSize = 2048; + static constexpr int64_t kOutputSize = 128; + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(cpu_backend != nullptr); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, kInputSize, kOutputSize); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, kInputSize, kOutputSize); + ggml_tensor * cpu_decode_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kInputSize, 1); + ggml_tensor * hrx_decode_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kInputSize, 1); + ggml_tensor * cpu_other_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kInputSize, 1); + ggml_tensor * hrx_other_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kInputSize, 1); + ggml_tensor * cpu_norm_weight = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, kInputSize); + ggml_tensor * hrx_norm_weight = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, kInputSize); + ggml_tensor * cpu_native_indices = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_native_indices = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_decode_rhs = build_rmsnorm_mul_graph(cpu_ctx, cpu_decode_input, cpu_norm_weight); + ggml_tensor * hrx_decode_rhs = build_rmsnorm_mul_graph(hrx_ctx, hrx_decode_input, hrx_norm_weight); + ggml_tensor * cpu_decode_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_decode_rhs); + ggml_tensor * hrx_decode_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_decode_rhs); + ggml_tensor * cpu_other_rhs = build_rmsnorm_mul_graph(cpu_ctx, cpu_other_input, cpu_norm_weight); + ggml_tensor * hrx_other_rhs = build_rmsnorm_mul_graph(hrx_ctx, hrx_other_input, hrx_norm_weight); + ggml_tensor * cpu_other_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_other_rhs); + ggml_tensor * hrx_other_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_other_rhs); + ggml_tensor * cpu_native_output = ggml_get_rows(cpu_ctx, cpu_weight, cpu_native_indices); + ggml_tensor * hrx_native_output = ggml_get_rows(hrx_ctx, hrx_weight, hrx_native_indices); + const bool cpu_host_source = source_kind == PackedWeightSourceKind::CpuHost; + const bool hrx_host_source = source_kind == PackedWeightSourceKind::HrxHost; + const bool coherent_subregion = source_kind == PackedWeightSourceKind::CoherentSubregion; + ggml_cgraph * cpu_decode = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_decode = ggml_new_graph(hrx_ctx); + ggml_cgraph * cpu_other = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_other = ggml_new_graph(hrx_ctx); + ggml_cgraph * cpu_native = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_native = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_decode != nullptr && hrx_decode != nullptr && cpu_other != nullptr && hrx_other != nullptr); + REQUIRE(cpu_native != nullptr && hrx_native != nullptr); + ggml_build_forward_expand(cpu_decode, cpu_decode_output); + ggml_build_forward_expand(hrx_decode, hrx_decode_output); + ggml_build_forward_expand(cpu_other, cpu_other_output); + ggml_build_forward_expand(hrx_other, hrx_other_output); + ggml_build_forward_expand(cpu_native, cpu_native_output); + ggml_build_forward_expand(hrx_native, hrx_native_output); + hrx_decode->uid = 951001; + hrx_native->uid = 951002; + hrx_other->uid = 951003; + + std::vector source = make_matmul_weight_bytes(weight_type, kInputSize, kOutputSize, 6); + ggml_backend_buffer_t imported = nullptr; + ggml_backend_buffer_t owner = nullptr; + ggml_tensor * owner_weight = nullptr; + uint8_t * coherent_owner_base = nullptr; + auto * backend_state = static_cast(hrx_backend->context); + if (cpu_host_source) { + imported = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), source.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, ggml_backend_buffer_get_base(imported)) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_set(hrx_weight, source.data(), 0, source.size()); + } else if (hrx_host_source) { + imported = ggml_backend_buft_alloc_buffer( + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(hrx_backend)), source.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, ggml_backend_buffer_get_base(imported)) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_set(hrx_weight, source.data(), 0, source.size()); + } else if (coherent_subregion) { + const bool original_direct = backend_state->device->use_direct_host_bindings; + backend_state->device->use_direct_host_bindings = true; + owner = ggml_backend_buft_alloc_buffer( + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(hrx_backend)), source.size() + 256); + backend_state->device->use_direct_host_bindings = original_direct; + REQUIRE(owner != nullptr); + auto * owner_base = static_cast(ggml_backend_buffer_get_base(owner)); + coherent_owner_base = owner_base; + owner_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, kInputSize, kOutputSize); + REQUIRE(ggml_backend_tensor_alloc(owner, owner_weight, owner_base + 256) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_set(owner_weight, source.data(), 0, source.size()); + imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(hrx_backend), owner_base + 256, source.size(), source.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, owner_base + 256) == GGML_STATUS_SUCCESS); + } else { + imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(hrx_backend), source.data(), source.size(), source.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, source.data()) == GGML_STATUS_SUCCESS); + } + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t decode_host = ggml_backend_buft_alloc_buffer( + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(hrx_backend)), 2 * ggml_nbytes(hrx_decode_input)); + ggml_backend_buffer_t other_host = ggml_backend_buft_alloc_buffer( + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(hrx_backend)), ggml_nbytes(hrx_other_input)); + REQUIRE(decode_host != nullptr && other_host != nullptr); + REQUIRE(ggml_backend_tensor_alloc(decode_host, hrx_decode_input, ggml_backend_buffer_get_base(decode_host)) == + GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(other_host, hrx_other_input, ggml_backend_buffer_get_base(other_host)) == + GGML_STATUS_SUCCESS); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + REQUIRE(hrx_buffer != nullptr && cpu_buffer != nullptr); + ggml_backend_tensor_set(cpu_weight, source.data(), 0, source.size()); + const std::vector decode_input = make_pattern_f32(kInputSize, 7, 0.01f); + const std::array native_indices = { 1, 17, 63 }; + set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, decode_input.data(), + decode_input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_other_input, hrx_backend, hrx_other_input, decode_input.data(), + decode_input.size() * sizeof(float)); + const std::vector norm_weight = make_weight(kInputSize); + set_tensor_pair_bytes(cpu_backend, cpu_norm_weight, hrx_backend, hrx_norm_weight, norm_weight.data(), + norm_weight.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_native_indices, hrx_backend, hrx_native_indices, native_indices.data(), + native_indices.size() * sizeof(int32_t)); + + auto compute_and_compare = [&](ggml_cgraph * cpu_graph, ggml_cgraph * hrx_graph, + ggml_tensor * cpu_output, ggml_tensor * hrx_output) { + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0f, 3.0e-2f); + }; + + if (source_kind != PackedWeightSourceKind::Imported) { + REQUIRE(!ggml_backend_hrx_weight_has_repackable_source(hrx_weight)); + ggml::hrx::ValueBufferBinding source_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(hrx_weight, source_binding)); + REQUIRE(source_binding.source_host_data != nullptr); + REQUIRE(source_binding.source_capacity >= source.size()); + REQUIRE((source_binding.buffer != nullptr) == coherent_subregion); + REQUIRE((source_binding.host_data != nullptr) == !coherent_subregion); + if (coherent_subregion) { + REQUIRE(source_binding.source_host_data == coherent_owner_base); + REQUIRE(source_binding.source_offset == 256); + REQUIRE(source_binding.source_capacity == source.size() + 256); + } + ggml_backend_hrx_cache_stats before_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_stats)); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats first_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &first_stats)); + const uint64_t expected_materializations = coherent_subregion ? 0 : 1; + REQUIRE(first_stats.weight_materializations == expected_materializations); + const size_t first_uploaded_bytes = first_stats.weight_uploaded_bytes; + + // Immutable GGUF weights retain one native copy while mutable inputs change. + const std::vector shifted_input = make_pattern_f32(kInputSize, 79, 0.06f); + set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, shifted_input.data(), + shifted_input.size() * sizeof(float)); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats repeated_stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &repeated_stats)); + REQUIRE(repeated_stats.graph_program_builds == first_stats.graph_program_builds); + REQUIRE(repeated_stats.prepared_program_builds == first_stats.prepared_program_builds); + REQUIRE(repeated_stats.prepared_program_hits > first_stats.prepared_program_hits); + REQUIRE(repeated_stats.weight_materializations == expected_materializations); + REQUIRE(repeated_stats.weight_uploaded_bytes == first_uploaded_bytes); + + auto lookup = backend_state->graph_programs.get_or_build( + *hrx_decode, ggml::hrx::get_qwen_kernel_corpus(), backend_state->device->architecture); + REQUIRE(lookup.valid()); + const auto * weight_value = lookup.program->graph().values().find_tensor(hrx_weight); + REQUIRE(weight_value != nullptr); + bool native_weight = false; + for (const auto & command : lookup.program->commands().commands) { + for (const auto & binding : command.bindings) { + if (binding.value == weight_value->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + native_weight = true; + } + } + } + REQUIRE(native_weight); + auto match = lookup.program->match_trusted_graph(*hrx_decode, true); + REQUIRE(match.valid()); + std::vector external; + for (const auto & slot : match.external_bindings) { + ggml::hrx::ValueBufferBinding value; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(slot.tensor, value)); + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = slot.tensor->ne[dimension]; + nb[dimension] = slot.tensor->nb[dimension]; + } + external.push_back(ggml::hrx::make_command_program_binding( + slot.value, value, slot.tensor->type, ne, nb, ggml_nbytes(slot.tensor) == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); + REQUIRE(bindings.valid()); + const auto * weight_binding = bindings.find(weight_value->id); + REQUIRE(weight_binding != nullptr); + REQUIRE(weight_binding->source_host_data == source_binding.source_host_data); + REQUIRE(weight_binding->source_offset == source_binding.source_offset); + REQUIRE(weight_binding->source_capacity == source_binding.source_capacity); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), + &corpus, &backend_state->kernel_executables, &backend_state->transient_arena, + &backend_state->host_transfers, &backend_state->host_weights, &backend_state->device->host_buffers, + &backend_state->graph_replay_state, + }; + auto prepared = ggml::hrx::prepare_command_program( + execution_context, lookup.program->commands(), bindings); + REQUIRE(prepared.valid()); + REQUIRE(prepared.resident_host_weights.size() == expected_materializations); + const bool staged_weight = std::any_of(prepared.host_staging.begin(), prepared.host_staging.end(), + [&](const auto & staging) { return staging.value == weight_value->id.value; }); + REQUIRE(!staged_weight); + compute_and_compare(cpu_other, hrx_other, cpu_other_output, hrx_other_output); + compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); + const std::vector first_rows = get_f32_tensor(cpu_backend, cpu_native_output); + const std::array changed_indices = { 2, 19, 71 }; + set_tensor_pair_bytes(cpu_backend, cpu_native_indices, hrx_backend, hrx_native_indices, + changed_indices.data(), sizeof(changed_indices)); + compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); + const std::vector changed_rows = get_f32_tensor(cpu_backend, cpu_native_output); + bool token_rows_changed = false; + for (size_t index = 0; index < changed_rows.size(); ++index) { + token_rows_changed = token_rows_changed || std::fabs(first_rows[index] - changed_rows[index]) > 1.0e-3f; + } + REQUIRE(token_rows_changed); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == expected_materializations); + REQUIRE(stats.weight_uploaded_bytes == first_uploaded_bytes); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_backend_buffer_free(other_host); + ggml_backend_buffer_free(decode_host); + ggml_backend_buffer_free(imported); + if (owner != nullptr) { + ggml_backend_buffer_free(owner); + } + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); + return; + } + + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_live_allocations == 0); + ggml_backend_hrx_imported_region input_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(hrx_decode_input, input_region)); + const uint64_t first_input_generation = input_region.generation; + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == 1); + REQUIRE(stats.weight_live_allocations == 1); + REQUIRE(stats.weight_live_bytes == stats.weight_uploaded_bytes); + const size_t first_weight_uploaded_bytes = stats.weight_uploaded_bytes; + const std::vector first_output = get_f32_tensor(cpu_backend, cpu_decode_output); + const std::vector changed_input = make_pattern_f32(kInputSize, 41, 0.04f); + set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, changed_input.data(), + changed_input.size() * sizeof(float)); + REQUIRE(ggml_backend_hrx_get_retained_region(hrx_decode_input, input_region)); + REQUIRE(input_region.generation != first_input_generation); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + const std::vector changed_output = get_f32_tensor(cpu_backend, cpu_decode_output); + bool output_changed = false; + for (size_t index = 0; index < changed_output.size(); ++index) { + const float stale_difference = std::fabs(first_output[index] - changed_output[index]); + const float allowed_difference = 1.0f + 3.0e-2f * std::fabs(changed_output[index]); + output_changed = output_changed || stale_difference > allowed_difference; + } + REQUIRE(output_changed); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.prepared_program_builds == 1); + REQUIRE(stats.prepared_program_hits >= 1); + REQUIRE(stats.weight_materializations == 1); + REQUIRE(stats.weight_uploaded_bytes == first_weight_uploaded_bytes); + auto lookup = backend_state->graph_programs.get_or_build( + *hrx_decode, ggml::hrx::get_qwen_kernel_corpus(), backend_state->device->architecture); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->imported_sources_current()); + auto bindings_for_match = [](const ggml::hrx::GraphProgramMatch & match) { + std::vector external; + for (const auto & slot : match.external_bindings) { + ggml::hrx::ValueBufferBinding value = {}; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(slot.tensor, value)); + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = slot.tensor->ne[dimension]; + nb[dimension] = slot.tensor->nb[dimension]; + } + external.push_back(ggml::hrx::make_command_program_binding( + slot.value, value, slot.tensor->type, ne, nb, ggml_nbytes(slot.tensor) == 0)); + } + return ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); + }; + const auto & replay_corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext replay_context = { + backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), + &replay_corpus, &backend_state->kernel_executables, &backend_state->transient_arena, + &backend_state->host_transfers, &backend_state->host_weights, &backend_state->device->host_buffers, + &backend_state->graph_replay_state, + }; + auto staging_match = lookup.program->match_host_staging_graph(*hrx_decode); + REQUIRE(staging_match.valid()); + auto staging_bindings = bindings_for_match(staging_match); + REQUIRE(staging_bindings.valid()); + { + const auto graph_replay = lookup.program->execute_with_result(replay_context, staging_bindings, false); + REQUIRE(graph_replay.success); + const char * replay_setting = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool replay_disabled = replay_setting != nullptr && replay_setting[0] != '\0' && replay_setting[0] != '0'; + const auto expected_graph_replay_event = replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : + ggml::hrx::HrxGraphReplayEvent::Hit; + REQUIRE(graph_replay.graph_replay_event == expected_graph_replay_event); + ggml_backend_synchronize(hrx_backend); + } + const std::string row64_layout = weight_type == GGML_TYPE_Q4_K ? + ggml::hrx::kQ4KPackedK256Row64Layout : ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + bool selected_row64 = false; + for (const ggml::hrx::Command & command : lookup.program->commands().commands) { + for (const ggml::hrx::CommandBinding & binding : command.bindings) { + selected_row64 = selected_row64 || binding.layout == row64_layout; + } + } + REQUIRE(selected_row64); + if (coherent_subregion) { + // This internal factory must preserve the owner CPU address even when the import binds directly by handle. + auto graph_import = ggml::hrx::import_ggml_graph(*hrx_decode); + REQUIRE(graph_import.valid()); + auto & values = graph_import.graph.values(); + for (ggml::hrx::ValueId id : values.external_value_ids()) { + const ggml::hrx::Value * value = values.find(id); + REQUIRE(value != nullptr); + REQUIRE(value->tensor != nullptr); + ggml::hrx::ValueBufferBinding source_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(value->tensor, source_binding)); + REQUIRE(values.bind_buffer(id, source_binding)); + } + const auto * weight_value = values.find_tensor(hrx_weight); + REQUIRE(weight_value != nullptr); + auto value_map_bindings = ggml::hrx::CommandProgramBindings::from_value_map(values); + REQUIRE(value_map_bindings.valid()); + const auto * weight_binding = value_map_bindings.find(weight_value->id); + REQUIRE(weight_binding != nullptr); + REQUIRE(weight_binding->buffer != nullptr); + REQUIRE(weight_binding->host_data == nullptr); + REQUIRE(weight_binding->source_host_data == coherent_owner_base); + REQUIRE(weight_binding->source_offset == 256); + REQUIRE(weight_binding->source_capacity == source.size() + 256); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), + &corpus, &backend_state->kernel_executables, &backend_state->transient_arena, + &backend_state->host_transfers, &backend_state->host_weights, &backend_state->device->host_buffers, + &backend_state->graph_replay_state, + }; + auto prepared = ggml::hrx::prepare_command_program( + execution_context, lookup.program->commands(), value_map_bindings); + REQUIRE(prepared.valid()); + REQUIRE(prepared.resident_host_weights.size() == 1); + ggml::hrx::RecordedCommandGraph recorded; + const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( + execution_context, lookup.program->commands(), value_map_bindings, prepared, recorded); + REQUIRE(launch.success); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_decode_output), + get_f32_tensor(cpu_backend, cpu_decode_output), 1.0f, 3.0e-2f); + } + compute_and_compare(cpu_other, hrx_other, cpu_other_output, hrx_other_output); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == 1); + REQUIRE(stats.weight_live_allocations == 1); + REQUIRE(stats.weight_reuses >= 1); + REQUIRE(stats.graph_program_builds == 1); + REQUIRE(stats.prepared_program_builds == 2); + const std::vector first_other_output = get_f32_tensor(cpu_backend, cpu_other_output); + const std::vector changed_other_input = make_pattern_f32(kInputSize, 67, 0.05f); + set_tensor_pair_bytes(cpu_backend, cpu_other_input, hrx_backend, hrx_other_input, changed_other_input.data(), + changed_other_input.size() * sizeof(float)); + const auto prepared_before_other = backend_state->prepared_programs.stats(); + compute_and_compare(cpu_other, hrx_other, cpu_other_output, hrx_other_output); + const std::vector changed_other_output = get_f32_tensor(cpu_backend, cpu_other_output); + bool other_output_changed = false; + for (size_t index = 0; index < changed_other_output.size(); ++index) { + const float stale_difference = std::fabs(first_other_output[index] - changed_other_output[index]); + const float allowed_difference = 1.0f + 3.0e-2f * std::fabs(changed_other_output[index]); + other_output_changed = other_output_changed || stale_difference > allowed_difference; + } + REQUIRE(other_output_changed); + const auto prepared_after_other = backend_state->prepared_programs.stats(); + REQUIRE(prepared_after_other.builds == prepared_before_other.builds); + REQUIRE(prepared_after_other.hits == prepared_before_other.hits + 1); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == 1); + auto full_match = lookup.program->match_trusted_graph(*hrx_other, true); + REQUIRE(full_match.valid()); + auto full_bindings = bindings_for_match(full_match); + REQUIRE(full_bindings.valid()); + { + const auto cache_replay = backend_state->prepared_programs.execute_with_result( + replay_context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), + full_bindings); + REQUIRE(cache_replay.success); + REQUIRE(cache_replay.graph_replay_event == ggml::hrx::HrxGraphReplayEvent::Hit); + ggml_backend_synchronize(hrx_backend); + } + std::vector readback(source.size()); + ggml_backend_tensor_get(hrx_weight, readback.data(), 0, readback.size()); + REQUIRE(readback == source); + + compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + const uint64_t native_materializations = coherent_subregion ? 1 : 2; + REQUIRE(stats.weight_materializations == native_materializations); + REQUIRE(stats.weight_live_allocations == (coherent_subregion ? 0 : 1)); + const uint64_t uploaded_bytes = stats.weight_uploaded_bytes; + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == 2 * native_materializations); + REQUIRE(stats.weight_live_allocations == (coherent_subregion ? 0 : 1)); + REQUIRE(stats.weight_uploaded_bytes == 2 * uploaded_bytes); + + // The preceding native consumer may retire packed residency; restore packed demand before testing input rebinding. + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats before_offset_change = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_offset_change)); + hrx_decode_input->data = static_cast(ggml_backend_buffer_get_base(decode_host)) + + ggml_nbytes(hrx_decode_input); + const std::vector shifted_input = make_pattern_f32(kInputSize, 79, 0.06f); + set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, shifted_input.data(), + shifted_input.size() * sizeof(float)); + REQUIRE(!lookup.program->imported_sources_current()); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats after_offset_change = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &after_offset_change)); + REQUIRE(after_offset_change.prepared_program_builds == before_offset_change.prepared_program_builds + 1); + REQUIRE(after_offset_change.weight_materializations == before_offset_change.weight_materializations); + + ggml_backend_buffer_t rebound_host = ggml_backend_buft_alloc_buffer( + ggml_backend_dev_host_buffer_type(ggml_backend_get_device(hrx_backend)), ggml_nbytes(hrx_decode_input)); + REQUIRE(rebound_host != nullptr); + hrx_decode_input->buffer = rebound_host; + hrx_decode_input->data = ggml_backend_buffer_get_base(rebound_host); + const std::vector rebound_input = make_pattern_f32(kInputSize, 83, 0.06f); + set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, rebound_input.data(), + rebound_input.size() * sizeof(float)); + REQUIRE(!lookup.program->imported_sources_current()); + ggml_backend_hrx_cache_stats before_rebind = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_rebind)); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + ggml_backend_hrx_cache_stats after_rebind = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &after_rebind)); + REQUIRE(after_rebind.graph_program_builds == before_rebind.graph_program_builds); + REQUIRE(after_rebind.prepared_program_builds == before_rebind.prepared_program_builds + 1); + REQUIRE(after_rebind.weight_materializations == before_rebind.weight_materializations); + + if (source_kind == PackedWeightSourceKind::Imported) { + // Exercise production preparation and recording: physical residency must remain counted after both + // the graph-program owner and weight lookup entries are gone, for as long as recording retains it. + auto full_match = lookup.program->match_trusted_graph(*hrx_decode, true); + REQUIRE(full_match.valid()); + const ggml::hrx::CommandProgram & commands = lookup.program->commands(); + std::vector external_bindings; + for (const auto & external : full_match.external_bindings) { + ggml::hrx::ValueBufferBinding value = {}; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, value)); + ggml::hrx::CommandProgramBinding binding = { + external.value, value.buffer, value.offset, value.length, value.identity, value.generation, + value.capacity, value.host_data, value.weight, ggml_nbytes(external.tensor) == 0, + }; + binding.source_host_data = value.host_data; + binding.source_offset = value.offset; + binding.source_capacity = value.capacity; + binding.source_type = external.tensor->type; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + binding.source_ne[dimension] = external.tensor->ne[dimension]; + binding.source_nb[dimension] = external.tensor->nb[dimension]; + } + external_bindings.push_back(binding); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); + REQUIRE(bindings.valid()); + ggml_backend_synchronize(hrx_backend); + backend_state->host_weights.clear(); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), + &corpus, &backend_state->kernel_executables, &backend_state->transient_arena, + &backend_state->host_transfers, &backend_state->host_weights, &backend_state->device->host_buffers, + &backend_state->graph_replay_state, + }; + auto prepared = ggml::hrx::prepare_command_program(execution_context, commands, bindings); + REQUIRE(prepared.valid()); + REQUIRE(prepared.resident_host_weights.size() == 1); + ggml::hrx::RecordedCommandGraph recorded; + const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( + execution_context, commands, bindings, prepared, recorded); + REQUIRE(launch.success); + REQUIRE(recorded.valid()); + ggml_backend_synchronize(hrx_backend); + prepared = {}; + backend_state->graph_programs.clear(); + backend_state->prepared_programs.clear(); + lookup = {}; + backend_state->weight_submissions.release_current_leases(); + backend_state->host_weights.clear(); + REQUIRE(backend_state->host_weights.stats().allocation_count == 1); + recorded = {}; + REQUIRE(backend_state->host_weights.stats().allocation_count == 0); + } + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_backend_buffer_free(other_host); + ggml_backend_buffer_free(decode_host); + ggml_backend_buffer_free(rebound_host); + ggml_backend_buffer_free(imported); + if (owner != nullptr) { + ggml_backend_buffer_free(owner); + } + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + +static void run_mixed_weight_layout_program_case(ggml_type weight_type) { + static constexpr int64_t input_size = 256; + static constexpr int64_t output_size = 64; + static constexpr float native_abs_tolerance = 2.0e-3f; + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(cpu != nullptr && backend != nullptr); + auto * state = static_cast(backend->context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_weight_second = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_weight_alias = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_weight_second = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * cpu_rows_second = ggml_get_rows(cpu_ctx, cpu_weight_second, cpu_ids); + ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_rows); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_rows); + ggml_tensor * hrx_rows_equivalent = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_output_equivalent = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_rows_equivalent); + ggml_tensor * hrx_rows_alias = ggml_get_rows(hrx_ctx, hrx_weight_alias, hrx_ids); + ggml_tensor * hrx_alias_output = ggml_add(hrx_ctx, hrx_rows, hrx_rows_alias); + ggml_tensor * hrx_rows_second = ggml_get_rows(hrx_ctx, hrx_weight_second, hrx_ids); + ggml_tensor * hrx_distinct_output = ggml_add(hrx_ctx, hrx_rows, hrx_rows_second); + ggml_tensor * cpu_double_rows = ggml_add(cpu_ctx, cpu_rows, cpu_rows); + ggml_tensor * cpu_distinct_rows = ggml_add(cpu_ctx, cpu_rows, cpu_rows_second); + ggml_tensor * hrx_weight_view = ggml_view_2d(hrx_ctx, hrx_weight, input_size, output_size, + hrx_weight->nb[1], 0); + ggml_tensor * hrx_view_rows = ggml_get_rows(hrx_ctx, hrx_weight_view, hrx_ids); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_graph_equivalent = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_graph_alias = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_graph_distinct = ggml_new_graph(hrx_ctx); + ggml_cgraph * hrx_graph_view = ggml_new_graph(hrx_ctx); + ggml_cgraph * cpu_double_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * cpu_distinct_graph = ggml_new_graph(cpu_ctx); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + ggml_build_forward_expand(hrx_graph_equivalent, hrx_output_equivalent); + ggml_build_forward_expand(hrx_graph_alias, hrx_alias_output); + ggml_build_forward_expand(hrx_graph_distinct, hrx_distinct_output); + ggml_build_forward_expand(hrx_graph_view, hrx_view_rows); + ggml_build_forward_expand(cpu_double_graph, cpu_double_rows); + ggml_build_forward_expand(cpu_distinct_graph, cpu_distinct_rows); + hrx_graph->uid = 951041; + hrx_graph_equivalent->uid = 951042; + hrx_graph_alias->uid = 951043; + hrx_graph_distinct->uid = 951044; + hrx_graph_view->uid = 951045; + + std::vector weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); + std::vector second_weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 11); + std::vector source = weight; + source.insert(source.end(), second_weight.begin(), second_weight.end()); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), source.data(), source.size(), source.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, source.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight_alias, source.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight_second, source.data() + weight.size()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); + ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + REQUIRE(hrx_tensors != nullptr && cpu_tensors != nullptr); + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + ggml_backend_tensor_set(cpu_weight_second, second_weight.data(), 0, second_weight.size()); + const std::array ids = { 0, 1, 63 }; + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); + + auto lookup = state->graph_programs.get_or_build( + *hrx_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(lookup.valid()); + auto imported_graph = ggml::hrx::import_ggml_graph(*hrx_graph); + REQUIRE(imported_graph.valid()); + const auto * weight_value = imported_graph.graph.values().find_tensor(hrx_weight); + REQUIRE(weight_value != nullptr); + const std::string row64 = weight_type == GGML_TYPE_Q4_K ? + ggml::hrx::kQ4KPackedK256Row64Layout : ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + bool native_selected = false; + bool row64_selected = false; + for (const auto & command : lookup.program->commands().commands) { + for (const auto & binding : command.bindings) { + if (binding.value == weight_value->id) { + native_selected |= binding.layout == ggml::hrx::kNativeWeightLayout; + row64_selected |= binding.layout == row64; + } + } + } + REQUIRE(native_selected && row64_selected); + + const auto before = state->host_weights.stats(); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_output), get_f32_tensor(cpu, cpu_output), 1.0f, 3.0e-2f); + const auto first = state->host_weights.stats(); + REQUIRE(first.misses == before.misses + 2); + REQUIRE(first.duplicate_warnings == before.duplicate_warnings + 1); + REQUIRE(first.allocation_count == 2); + REQUIRE(first.resident_bytes == 2 * weight.size()); + REQUIRE(first.peak_resident_bytes == first.resident_bytes); + ggml_backend_hrx_cache_stats public_first = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &public_first)); + REQUIRE(public_first.weight_duplicate_warnings == first.duplicate_warnings); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_output), get_f32_tensor(cpu, cpu_output), 1.0f, 3.0e-2f); + const auto repeated = state->host_weights.stats(); + REQUIRE(repeated.misses == first.misses); + REQUIRE(repeated.duplicate_warnings == first.duplicate_warnings); + const auto repeated_graph_stats = state->graph_programs.stats(); + REQUIRE(repeated_graph_stats.prepared_program_hits > 0); + REQUIRE(repeated_graph_stats.weight_demand_builds >= 1); + REQUIRE(repeated_graph_stats.weight_demand_hits >= 1); + ggml_backend_hrx_cache_stats public_repeated = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &public_repeated)); + REQUIRE(public_repeated.weight_duplicate_warnings == public_first.weight_duplicate_warnings); + ggml_backend_hrx_cache_stats cache_before_equivalent = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_before_equivalent)); + const auto prepared_cache_before_equivalent = state->prepared_programs.stats(); + const auto graph_cache_before_equivalent = state->graph_programs.stats(); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph_equivalent) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_output_equivalent), get_f32_tensor(cpu, cpu_output), 1.0f, 3.0e-2f); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph_equivalent) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const auto equivalent = state->host_weights.stats(); + REQUIRE(equivalent.misses == first.misses); + REQUIRE(equivalent.duplicate_warnings == first.duplicate_warnings); + ggml_backend_hrx_cache_stats cache_after_equivalent = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_after_equivalent)); + REQUIRE(cache_after_equivalent.prepared_program_hits > cache_before_equivalent.prepared_program_hits); + REQUIRE(state->graph_programs.stats().builds == graph_cache_before_equivalent.builds); + REQUIRE(state->prepared_programs.stats().builds == prepared_cache_before_equivalent.builds + 1); + REQUIRE(state->prepared_programs.stats().hits > prepared_cache_before_equivalent.hits); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_double_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph_alias) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const auto cpu_double_reference = get_f32_tensor(cpu, cpu_double_rows); + require_close(get_f32_tensor(backend, hrx_alias_output), cpu_double_reference, native_abs_tolerance); + const auto aliases = state->host_weights.stats(); + REQUIRE(aliases.misses == first.misses); + REQUIRE(aliases.duplicate_warnings == first.duplicate_warnings); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph_view) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_view_rows), get_f32_tensor(cpu, cpu_rows), native_abs_tolerance); + REQUIRE(state->host_weights.stats().misses == aliases.misses); + REQUIRE(state->host_weights.stats().duplicate_warnings == aliases.duplicate_warnings); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_distinct_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph_distinct) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const auto cpu_distinct_reference = get_f32_tensor(cpu, cpu_distinct_rows); + float distinct_reference_difference = 0.0f; + for (size_t i = 0; i < cpu_distinct_reference.size(); ++i) { + distinct_reference_difference = std::max( + distinct_reference_difference, std::fabs(cpu_distinct_reference[i] - cpu_double_reference[i])); + } + REQUIRE(distinct_reference_difference > native_abs_tolerance); + require_close(get_f32_tensor(backend, hrx_distinct_output), cpu_distinct_reference, native_abs_tolerance); + const auto distinct = state->host_weights.stats(); + REQUIRE(distinct.misses == aliases.misses + 1); + REQUIRE(distinct.duplicate_warnings == aliases.duplicate_warnings); + + // Existing selected kernels also exercise an initialization-only native weight request. + auto & values = imported_graph.graph.values(); + for (const auto id : values.external_value_ids()) { + const auto * value = values.find(id); + REQUIRE(value != nullptr && value->tensor != nullptr); + ggml::hrx::ValueBufferBinding buffer; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(value->tensor, buffer)); + REQUIRE(values.bind_buffer(id, buffer)); + } + const auto bindings = ggml::hrx::CommandProgramBindings::from_value_map(values); + REQUIRE(bindings.valid()); + const auto * ids_value = values.find_tensor(hrx_ids); + REQUIRE(ids_value != nullptr); + const auto * ids_binding = bindings.find(ids_value->id); + REQUIRE(ids_binding != nullptr && !ids_binding->weight); + const auto & selected = lookup.program->commands(); + bool found_packed_weight_request = false; + bool found_native_ids_request = false; + for (const auto & command : selected.commands) { + for (const auto & request : command.bindings) { + if (request.value == ids_value->id && request.layout == ggml::hrx::kNativeWeightLayout) { + found_native_ids_request = true; + REQUIRE(bindings.find(request) == ids_binding); + } + if (request.value != weight_value->id || request.layout != row64) { + continue; + } + found_packed_weight_request = true; + REQUIRE(bindings.find(request) == nullptr); + auto nonweight_packed_request = request; + nonweight_packed_request.value = ids_value->id; + REQUIRE(bindings.find(nonweight_packed_request) == nullptr); + } + } + REQUIRE(found_packed_weight_request); + REQUIRE(found_native_ids_request); + ggml::hrx::CommandProgram constructed; + constructed.initialization_commands = selected.initialization_commands; + constructed.commands = selected.commands; + constructed.transients = selected.transients; + constructed.completion_counters = selected.completion_counters; + constructed.constant_initializations = selected.constant_initializations; + constructed.activation_publication_diagnostics = selected.activation_publication_diagnostics; + REQUIRE(!constructed.commands.empty()); + bool first_reads_native_weight = false; + for (const auto & binding : constructed.commands.front().bindings) { + first_reads_native_weight |= binding.value == weight_value->id && + binding.layout == ggml::hrx::kNativeWeightLayout; + } + REQUIRE(first_reads_native_weight); + constructed.initialization_commands.push_back(constructed.commands.front()); + constructed.commands.erase(constructed.commands.begin()); + for (auto & command : constructed.commands) { + --command.ordinal; + command.dependencies.erase(std::remove(command.dependencies.begin(), command.dependencies.end(), 0), + command.dependencies.end()); + for (auto & dependency : command.dependencies) { + --dependency; + } + } + const auto demands = ggml::hrx::collect_command_program_weight_demands(constructed, bindings); + REQUIRE(demands.valid()); + REQUIRE(demands.weights.size() == 2); + const std::array init_ids = { 2, 7, 61 }; + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, init_ids.data(), sizeof(init_ids)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + const std::vector poison_rows(static_cast(ggml_nelements(hrx_rows)), 1000.0f); + const std::vector poison_output(static_cast(ggml_nelements(hrx_output)), -1000.0f); + ggml_backend_tensor_set(hrx_rows, poison_rows.data(), 0, poison_rows.size() * sizeof(float)); + ggml_backend_tensor_set(hrx_output, poison_output.data(), 0, poison_output.size() * sizeof(float)); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + state->device->device, state->stream, state->device->architecture.c_str(), &corpus, + &state->kernel_executables, &state->transient_arena, &state->host_transfers, + &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, + }; + auto prepared = ggml::hrx::prepare_command_program(execution_context, constructed, bindings); + REQUIRE(prepared.valid()); + hrx_buffer_t native_buffer = nullptr; + hrx_buffer_t row64_buffer = nullptr; + for (const auto & command : prepared.initialization_commands) { + for (const auto & binding : command.kernel.bindings) { + if (binding.binding.value == weight_value->id && binding.binding.layout == ggml::hrx::kNativeWeightLayout) { + native_buffer = binding.ref.buffer; + REQUIRE(binding.ref.length == binding.binding.length); + } + } + } + for (const auto & command : prepared.commands) { + for (const auto & binding : command.kernel.bindings) { + if (binding.binding.value == weight_value->id && binding.binding.layout == row64) { + row64_buffer = binding.ref.buffer; + REQUIRE(binding.ref.length == binding.binding.length); + } + } + } + REQUIRE(native_buffer != nullptr && row64_buffer != nullptr && native_buffer != row64_buffer); + REQUIRE(ggml::hrx::bind_and_execute_prepared_command_program(execution_context, constructed, bindings, prepared)); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_output), get_f32_tensor(cpu, cpu_output), 1.0f, 3.0e-2f); + + const std::array replay_ids = { 3, 9, 59 }; + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, replay_ids.data(), sizeof(replay_ids)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_set(hrx_rows, poison_rows.data(), 0, poison_rows.size() * sizeof(float)); + ggml_backend_tensor_set(hrx_output, poison_output.data(), 0, poison_output.size() * sizeof(float)); + ggml::hrx::RecordedCommandGraph recorded; + const auto replay = ggml::hrx::bind_and_launch_recorded_command_graph( + execution_context, constructed, bindings, prepared, recorded); + REQUIRE(replay.success); + REQUIRE(recorded.valid()); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_output), get_f32_tensor(cpu, cpu_output), 1.0f, 3.0e-2f); + + std::vector readback(weight.size()); + ggml_backend_tensor_get(hrx_weight, readback.data(), 0, readback.size()); + REQUIRE(readback == weight); + + ggml_backend_buffer_free(hrx_tensors); + ggml_backend_buffer_free(cpu_tensors); + ggml_backend_buffer_free(imported); + ggml_free(hrx_ctx); + ggml_free(cpu_ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); +} + +static void run_native_cpu_weight_during_import_layout_transition_case(ggml_type weight_type) { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(cpu != nullptr && backend != nullptr); + auto * state = static_cast(backend->context); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + + ggml_tensor * cpu_host_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_imported_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * hrx_host_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_imported_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); + ggml_tensor * cpu_host_rows = ggml_get_rows(cpu_ctx, cpu_host_weight, cpu_ids); + ggml_tensor * hrx_host_rows = ggml_get_rows(hrx_ctx, hrx_host_weight, hrx_ids); + ggml_tensor * cpu_imported_rows = ggml_get_rows(cpu_ctx, cpu_imported_weight, cpu_ids); + ggml_tensor * hrx_imported_rows = ggml_get_rows(hrx_ctx, hrx_imported_weight, hrx_ids); + ggml_tensor * cpu_imported_matmul = ggml_mul_mat(cpu_ctx, cpu_imported_weight, cpu_input); + ggml_tensor * hrx_imported_matmul = ggml_mul_mat(hrx_ctx, hrx_imported_weight, hrx_input); + ggml_cgraph * cpu_packed = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_packed = ggml_new_graph(hrx_ctx); + ggml_cgraph * cpu_native = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_native = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_packed != nullptr && hrx_packed != nullptr && cpu_native != nullptr && hrx_native != nullptr); + ggml_build_forward_expand(cpu_packed, cpu_host_rows); + ggml_build_forward_expand(cpu_packed, cpu_imported_matmul); + ggml_build_forward_expand(hrx_packed, hrx_host_rows); + ggml_build_forward_expand(hrx_packed, hrx_imported_matmul); + ggml_build_forward_expand(cpu_native, cpu_host_rows); + ggml_build_forward_expand(cpu_native, cpu_imported_rows); + ggml_build_forward_expand(hrx_native, hrx_host_rows); + ggml_build_forward_expand(hrx_native, hrx_imported_rows); + hrx_packed->uid = 951061; + hrx_native->uid = 951062; + + std::vector host_bytes = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); + std::vector imported_bytes = make_matmul_weight_bytes(weight_type, input_size, output_size, 11); + ggml_backend_buffer_t host_buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), host_bytes.size()); + REQUIRE(host_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(host_buffer, hrx_host_weight, ggml_backend_buffer_get_base(host_buffer)) == + GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(host_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_tensor_set(hrx_host_weight, host_bytes.data(), 0, host_bytes.size()); + ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), imported_bytes.data(), imported_bytes.size(), imported_bytes.size()); + REQUIRE(imported != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported, hrx_imported_weight, imported_bytes.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); + ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + REQUIRE(hrx_tensors != nullptr && cpu_tensors != nullptr); + ggml_backend_tensor_set(cpu_host_weight, host_bytes.data(), 0, host_bytes.size()); + ggml_backend_tensor_set(cpu_imported_weight, imported_bytes.data(), 0, imported_bytes.size()); + const std::array ids = { 1, 17, 63 }; + const std::vector input = make_pattern_f32(input_size * 3, 7, 0.01f); + set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); + set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); + + const std::string packed_layout = weight_type == GGML_TYPE_Q4_K ? + ggml::hrx::kQ4KPackedK256Row64Layout : ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + auto check_layouts = [&](ggml_cgraph * graph, bool expect_packed) { + auto lookup = state->graph_programs.get_or_build( + *graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(lookup.valid()); + const auto * host_value = lookup.program->graph().values().find_tensor(hrx_host_weight); + const auto * imported_value = lookup.program->graph().values().find_tensor(hrx_imported_weight); + REQUIRE(host_value != nullptr && imported_value != nullptr); + bool host_native = false; + bool imported_selected = false; + for (const auto & command : lookup.program->commands().commands) { + for (const auto & binding : command.bindings) { + if (binding.value == host_value->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + host_native = true; + } + if (binding.value == imported_value->id) { + const std::string & expected = expect_packed ? packed_layout : ggml::hrx::kNativeWeightLayout; + imported_selected = imported_selected || binding.layout == expected; + } + } + } + REQUIRE(host_native && imported_selected); + }; + check_layouts(hrx_packed, true); + check_layouts(hrx_native, false); + + auto compute_and_compare = [&](ggml_cgraph * cpu_graph, ggml_cgraph * hrx_graph, + ggml_tensor * cpu_imported_output, ggml_tensor * hrx_imported_output) { + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, hrx_host_rows), get_f32_tensor(cpu, cpu_host_rows), 2.0e-3f); + require_close(get_f32_tensor(backend, hrx_imported_output), + get_f32_tensor(cpu, cpu_imported_output), 1.0f, 3.0e-2f); + }; + compute_and_compare(cpu_packed, hrx_packed, cpu_imported_matmul, hrx_imported_matmul); + auto first = state->host_weights.stats(); + REQUIRE(first.misses == 2); + REQUIRE(first.allocation_count == 2); + compute_and_compare(cpu_native, hrx_native, cpu_imported_rows, hrx_imported_rows); + auto native = state->host_weights.stats(); + REQUIRE(native.misses == first.misses + 1); + REQUIRE(native.allocation_count == 2); + compute_and_compare(cpu_packed, hrx_packed, cpu_imported_matmul, hrx_imported_matmul); + auto packed_again = state->host_weights.stats(); + REQUIRE(packed_again.misses == native.misses + 1); + REQUIRE(packed_again.allocation_count == 2); + + ggml_backend_buffer_free(hrx_tensors); + ggml_backend_buffer_free(cpu_tensors); + ggml_backend_buffer_free(imported); + ggml_backend_buffer_free(host_buffer); + ggml_free(hrx_ctx); + ggml_free(cpu_ctx); + ggml_backend_free(backend); + ggml_backend_free(cpu); +} + +static void run_q4_prefill_source_eligibility_case() { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + constexpr int64_t token_count = 256; + ggml_backend_t backend = ggml_backend_hrx_init(0); + ggml_backend_t cpu = init_cpu_backend(); + REQUIRE(backend != nullptr); + auto * backend_state = static_cast(backend->context); + const std::vector weight_bytes = + make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 17); + const std::vector input_bytes = make_pattern_f32(input_size * token_count, 23, 0.01f); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); + REQUIRE(cpu_buffer != nullptr); + ggml_backend_tensor_set(cpu_weight, weight_bytes.data(), 0, weight_bytes.size()); + ggml_backend_tensor_set(cpu_input, input_bytes.data(), 0, input_bytes.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); + const std::vector expected = get_f32_tensor(cpu, cpu_output); + + struct Fixture { + ggml_context * ctx = nullptr; + ggml_tensor * weight = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t weight_buffer = nullptr; + ggml_backend_buffer_t other_buffer = nullptr; + std::vector source; + }; + std::vector fixtures; + fixtures.reserve(5); + auto make_fixture = [&](bool imported, uint64_t uid) { + Fixture fixture; + fixture.ctx = ggml_init(params); + REQUIRE(fixture.ctx != nullptr); + fixture.weight = ggml_new_tensor_2d(fixture.ctx, GGML_TYPE_Q4_K, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(fixture.ctx, GGML_TYPE_F32, input_size, token_count); + fixture.output = ggml_mul_mat(fixture.ctx, fixture.weight, input); + fixture.graph = ggml_new_graph(fixture.ctx); + ggml_build_forward_expand(fixture.graph, fixture.output); + fixture.graph->uid = uid; + fixture.source = weight_bytes; + if (imported) { + fixture.weight_buffer = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), fixture.source.data(), fixture.source.size(), fixture.source.size()); + REQUIRE(fixture.weight_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(fixture.weight_buffer, fixture.weight, fixture.source.data()) == + GGML_STATUS_SUCCESS); + } else { + const auto buft = ggml_backend_get_default_buffer_type(backend); + const size_t alignment = ggml_backend_buft_get_alignment(buft); + const size_t bytes = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, fixture.weight), alignment); + fixture.weight_buffer = ggml_backend_buft_alloc_buffer(buft, bytes); + REQUIRE(fixture.weight_buffer != nullptr); + ggml_tallocr allocator = ggml_tallocr_new(fixture.weight_buffer); + REQUIRE(ggml_tallocr_alloc(&allocator, fixture.weight) == GGML_STATUS_SUCCESS); + } + ggml_backend_buffer_set_usage(fixture.weight_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + fixture.other_buffer = ggml_backend_alloc_ctx_tensors(fixture.ctx, backend); + REQUIRE(fixture.other_buffer != nullptr); + if (!imported) { + std::vector upload_staging = weight_bytes; + ggml_backend_tensor_set(fixture.weight, upload_staging.data(), 0, upload_staging.size()); + std::fill(upload_staging.begin(), upload_staging.end(), 0); + } + ggml_backend_tensor_set(input, input_bytes.data(), 0, input_bytes.size() * sizeof(float)); + return fixture; + }; + fixtures.push_back(make_fixture(true, 960001)); + fixtures.push_back(make_fixture(false, 960001)); + fixtures.push_back(make_fixture(false, 960001)); + fixtures.push_back(make_fixture(true, 960003)); + fixtures.push_back(make_fixture(true, 960004)); + + auto verify_and_run = [&](Fixture & fixture, bool imported, const std::vector & expected_weight, + const std::vector & expected_output) { + REQUIRE(ggml_backend_hrx_weight_has_repackable_source(fixture.weight) == imported); + ggml::hrx::GraphImportResult graph_import = ggml::hrx::import_ggml_graph(*fixture.graph); + REQUIRE(graph_import.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(graph_import.graph, { "gfx1151" })); + const ggml::hrx::Value * value = graph_import.graph.values().find_tensor(fixture.weight); + REQUIRE(value != nullptr); + bool found_weight = false; + bool found_expected_kernel = false; + for (const auto & dispatch : scheduler.plan().dispatches) { + const std::string name = kernel_name_for_id(dispatch.kernel.kernel_id); + const char * expected_kernel = imported ? "loom_libs:ggml_mul_mat_q4_k_f16_wmma_prefill_wave32" : + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256"; + found_expected_kernel |= name == expected_kernel; + for (const auto & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == (imported ? ggml::hrx::kQ4KPackedK256Row64Layout : + ggml::hrx::kNativeWeightLayout)); + found_weight = true; + } + } + } + REQUIRE(found_weight); + REQUIRE(found_expected_kernel); + const auto before = backend_state->host_weights.stats(); + REQUIRE(ggml_backend_graph_compute(backend, fixture.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + if (!imported) { + const auto after = backend_state->host_weights.stats(); + REQUIRE(after.misses == before.misses); + REQUIRE(after.allocation_count == before.allocation_count); + } + std::vector readback(expected_weight.size()); + ggml_backend_tensor_get(fixture.weight, readback.data(), 0, readback.size()); + REQUIRE(readback == expected_weight); + require_close(get_f32_tensor(backend, fixture.output), expected_output, 1.0f, 3.0e-2f); + }; + for (size_t index = 0; index < fixtures.size(); ++index) { + const bool imported = index == 0 || index >= 3; + verify_and_run(fixtures[index], imported, weight_bytes, expected); + } + ggml_backend_hrx_cache_stats cache = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache)); + REQUIRE(cache.graph_program_builds == 3); + REQUIRE(cache.graph_program_hits >= 2); + REQUIRE(cache.prepared_program_builds >= 3); + verify_and_run(fixtures[4], true, weight_bytes, expected); + ggml_backend_hrx_cache_stats structural_replay = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &structural_replay)); + REQUIRE(structural_replay.graph_program_builds == cache.graph_program_builds); + REQUIRE(structural_replay.prepared_program_hits > cache.prepared_program_hits); + + // Keep the graph and node pointers fixed while changing only the live weight binding. + Fixture & replay = fixtures[0]; + ggml_backend_buffer_t imported_buffer = replay.weight->buffer; + void * imported_data = replay.weight->data; + for (const bool imported : { true, false, true }) { + replay.weight->buffer = imported ? imported_buffer : fixtures[1].weight->buffer; + replay.weight->data = imported ? imported_data : fixtures[1].weight->data; + ggml_backend_hrx_cache_stats before = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &before)); + verify_and_run(replay, imported, weight_bytes, expected); + verify_and_run(replay, imported, weight_bytes, expected); + ggml_backend_hrx_cache_stats after = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &after)); + REQUIRE(after.graph_program_hits > before.graph_program_hits); + REQUIRE(after.prepared_program_hits > before.prepared_program_hits); + } + + for (Fixture & fixture : fixtures) { + ggml_backend_buffer_free(fixture.other_buffer); + ggml_backend_buffer_free(fixture.weight_buffer); + ggml_free(fixture.ctx); + } + ggml_backend_buffer_free(cpu_buffer); + ggml_free(cpu_ctx); + ggml_backend_free(cpu); + ggml_backend_free(backend); +} + static void run_router_top8_case(int64_t token_count) { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -4201,6 +6414,16 @@ static auto allocate_test_weight(ggml_backend_t backend, ggml_tensor * weight) { return buffer; } +static auto import_test_weight(ggml_backend_t backend, ggml_tensor * weight, std::vector & bytes) { + auto buffer = std::unique_ptr( + ggml_backend_dev_buffer_from_host_ptr(ggml_backend_get_device(backend), bytes.data(), bytes.size(), + bytes.size()), ggml_backend_buffer_free); + REQUIRE(buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(buffer.get(), weight, bytes.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(buffer.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + return buffer; +} + static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, @@ -4208,7 +6431,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, int64_t input_size = kQwenHiddenSize, bool normalize_input = false, bool square_output = false, - ggml_type cache_type = GGML_TYPE_COUNT) { + ggml_type cache_type = GGML_TYPE_COUNT, + bool verify_uploaded_native = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -4279,6 +6503,9 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); + if (verify_uploaded_native) { + require_native_weight_bindings(hrx_graph, { hrx_weight }); + } const std::vector weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); const bool weight_is_f16 = weight_type == GGML_TYPE_F16; @@ -4294,7 +6521,19 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, input[i] = static_cast(value) * 0.0078125f; } } - set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + if (verify_uploaded_native) { + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + std::vector upload_staging = weight; + ggml_backend_tensor_set(hrx_weight, upload_staging.data(), 0, upload_staging.size()); + std::fill(upload_staging.begin(), upload_staging.end(), 0); + upload_staging.clear(); + upload_staging.shrink_to_fit(); + std::vector readback(weight.size()); + ggml_backend_tensor_get(hrx_weight, readback.data(), 0, readback.size()); + REQUIRE(readback == weight); + } else { + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + } set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); if (normalize_input) { const std::vector norm_w = make_weight(input_size); @@ -4315,6 +6554,11 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const float abs_tolerance = cache_type != GGML_TYPE_COUNT ? 2.0e-3f : 1.0f; require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), abs_tolerance, cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); + if (verify_uploaded_native) { + std::vector readback(weight.size()); + ggml_backend_tensor_get(hrx_weight, readback.data(), 0, readback.size()); + REQUIRE(readback == weight); + } ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -4495,7 +6739,9 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ int64_t token_count, int64_t output_size, ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, - int64_t input_size = kQwenHiddenSize) { + int64_t input_size = kQwenHiddenSize, + bool verify_uploaded_native = false, + bool import_weights = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -4523,8 +6769,12 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); - auto hrx_gate_buffer = allocate_test_weight(hrx_backend, hrx_gate_weight); - auto hrx_up_buffer = allocate_test_weight(hrx_backend, hrx_up_weight); + std::vector gate_weight = make_matmul_weight_bytes(gate_weight_type, input_size, output_size, 6); + std::vector up_weight = make_matmul_weight_bytes(up_weight_type, input_size, output_size, 11); + auto hrx_gate_buffer = import_weights ? import_test_weight(hrx_backend, hrx_gate_weight, gate_weight) : + allocate_test_weight(hrx_backend, hrx_gate_weight); + auto hrx_up_buffer = import_weights ? import_test_weight(hrx_backend, hrx_up_weight, up_weight) : + allocate_test_weight(hrx_backend, hrx_up_weight); ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); @@ -4539,20 +6789,51 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); + if (verify_uploaded_native) { + require_native_weight_bindings(hrx_graph, { hrx_gate_weight, hrx_up_weight }); + } + if (import_weights) { + require_packed_weight_bindings(hrx_graph, { hrx_gate_weight, hrx_up_weight }); + } - const std::vector gate_weight = make_matmul_weight_bytes(gate_weight_type, input_size, output_size, 6); - const std::vector up_weight = make_matmul_weight_bytes(up_weight_type, input_size, output_size, 11); const std::vector input(static_cast(input_size * token_count), 0.00390625f); - set_tensor_pair_bytes(cpu_backend, cpu_gate_weight, hrx_backend, hrx_gate_weight, gate_weight.data(), - gate_weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_up_weight, hrx_backend, hrx_up_weight, up_weight.data(), up_weight.size()); + if (verify_uploaded_native) { + ggml_backend_tensor_set(cpu_gate_weight, gate_weight.data(), 0, gate_weight.size()); + ggml_backend_tensor_set(cpu_up_weight, up_weight.data(), 0, up_weight.size()); + std::vector gate_staging = gate_weight; + std::vector up_staging = up_weight; + ggml_backend_tensor_set(hrx_gate_weight, gate_staging.data(), 0, gate_staging.size()); + ggml_backend_tensor_set(hrx_up_weight, up_staging.data(), 0, up_staging.size()); + std::fill(gate_staging.begin(), gate_staging.end(), 0); + std::fill(up_staging.begin(), up_staging.end(), 0); + } else if (import_weights) { + ggml_backend_tensor_set(cpu_gate_weight, gate_weight.data(), 0, gate_weight.size()); + ggml_backend_tensor_set(cpu_up_weight, up_weight.data(), 0, up_weight.size()); + } else { + set_tensor_pair_bytes(cpu_backend, cpu_gate_weight, hrx_backend, hrx_gate_weight, gate_weight.data(), + gate_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_up_weight, hrx_backend, hrx_up_weight, up_weight.data(), up_weight.size()); + } set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + auto * backend_state = static_cast(hrx_backend->context); + const auto before = backend_state->host_weights.stats(); REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0e-4f, 1.0e-3f); + if (verify_uploaded_native) { + const auto after = backend_state->host_weights.stats(); + REQUIRE(after.misses == before.misses); + REQUIRE(after.allocation_count == before.allocation_count); + std::vector gate_readback(gate_weight.size()); + std::vector up_readback(up_weight.size()); + ggml_backend_tensor_get(hrx_gate_weight, gate_readback.data(), 0, gate_readback.size()); + ggml_backend_tensor_get(hrx_up_weight, up_readback.data(), 0, up_readback.size()); + REQUIRE(gate_readback == gate_weight); + REQUIRE(up_readback == up_weight); + } ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -4562,6 +6843,73 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ ggml_backend_free(hrx_backend); } +static void run_mixed_symmetric_source_eligibility_case() { + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 64; + constexpr int64_t token_count = 2; + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, input_size, output_size); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + std::vector gate_bytes = make_matmul_weight_bytes(GGML_TYPE_Q5_K, input_size, output_size, 6); + std::vector up_bytes = make_matmul_weight_bytes(GGML_TYPE_IQ4_XS, input_size, output_size, 11); + ggml_backend_buffer_t gate_buffer = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), gate_bytes.data(), gate_bytes.size(), gate_bytes.size()); + ggml_backend_buffer_t up_buffer = ggml_backend_dev_buffer_from_host_ptr( + ggml_backend_get_device(backend), up_bytes.data(), up_bytes.size(), up_bytes.size()); + REQUIRE(gate_buffer != nullptr); + REQUIRE(up_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(gate_buffer, gate_weight, gate_bytes.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(up_buffer, up_weight, up_bytes.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(gate_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_set_usage(up_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t other_buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(other_buffer != nullptr); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + const ggml::hrx::Value * gate_value = imported.graph.values().find_tensor(gate_weight); + const ggml::hrx::Value * up_value = imported.graph.values().find_tensor(up_weight); + REQUIRE(gate_value != nullptr); + REQUIRE(up_value != nullptr); + bool found_kernel = false; + size_t converted_bindings = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + found_kernel |= kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma"; + for (const auto & binding : dispatch.bindings) { + const bool is_weight = binding.value == gate_value->id || binding.value == up_value->id; + if (is_weight) { + REQUIRE(binding.layout == ggml::hrx::kSymmetricI4K32EightGroupsShared4Layout); + ++converted_bindings; + } + } + } + REQUIRE(found_kernel); + REQUIRE(converted_bindings == 2); + ggml_backend_buffer_free(other_buffer); + ggml_backend_buffer_free(up_buffer); + ggml_backend_buffer_free(gate_buffer); + ggml_free(ctx); + ggml_backend_free(backend); + + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q5_K, GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", + token_count, output_size, GGML_GLU_OP_SWIGLU, input_size, true); +} + static void run_dense_matmul_packed_glu_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, @@ -4629,7 +6977,8 @@ static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_t ggml_type rhs_weight_type, const char * expected_kernel, int64_t token_count, - int64_t output_size) { + int64_t output_size, + bool import_weights = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -4658,8 +7007,12 @@ static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_t REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); - auto hrx_lhs_buffer = allocate_test_weight(hrx_backend, hrx_lhs_weight); - auto hrx_rhs_buffer = allocate_test_weight(hrx_backend, hrx_rhs_weight); + std::vector lhs_weight = make_matmul_weight_bytes(lhs_weight_type, input_size, output_size, 6); + std::vector rhs_weight = make_matmul_weight_bytes(rhs_weight_type, input_size, output_size, 11); + auto hrx_lhs_buffer = import_weights ? import_test_weight(hrx_backend, hrx_lhs_weight, lhs_weight) : + allocate_test_weight(hrx_backend, hrx_lhs_weight); + auto hrx_rhs_buffer = import_weights ? import_test_weight(hrx_backend, hrx_rhs_weight, rhs_weight) : + allocate_test_weight(hrx_backend, hrx_rhs_weight); ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); @@ -4674,14 +7027,20 @@ static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_t ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); + if (import_weights) { + require_packed_weight_bindings(hrx_graph, { hrx_lhs_weight, hrx_rhs_weight }); + } - const std::vector lhs_weight = make_matmul_weight_bytes(lhs_weight_type, input_size, output_size, 6); - const std::vector rhs_weight = make_matmul_weight_bytes(rhs_weight_type, input_size, output_size, 11); const std::vector input(static_cast(input_size * token_count), 0.00390625f); - set_tensor_pair_bytes(cpu_backend, cpu_lhs_weight, hrx_backend, hrx_lhs_weight, lhs_weight.data(), - lhs_weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_rhs_weight, hrx_backend, hrx_rhs_weight, rhs_weight.data(), - rhs_weight.size()); + if (import_weights) { + ggml_backend_tensor_set(cpu_lhs_weight, lhs_weight.data(), 0, lhs_weight.size()); + ggml_backend_tensor_set(cpu_rhs_weight, rhs_weight.data(), 0, rhs_weight.size()); + } else { + set_tensor_pair_bytes(cpu_backend, cpu_lhs_weight, hrx_backend, hrx_lhs_weight, lhs_weight.data(), + lhs_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_rhs_weight, hrx_backend, hrx_rhs_weight, rhs_weight.data(), + rhs_weight.size()); + } set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); @@ -4729,8 +7088,17 @@ static void run_vector_q8_publish_cpu_reference_case() { REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); - auto hrx_first_weight_buffer = allocate_test_weight(hrx_backend, hrx_first_weight); - auto hrx_second_weight_buffer = allocate_test_weight(hrx_backend, hrx_second_weight); + std::vector second_weight = + make_matmul_weight_bytes(GGML_TYPE_Q4_K, intermediate_size, output_size, 11); + auto hrx_first_weight_buffer = allocate_test_weight(hrx_backend, hrx_first_weight); + auto hrx_second_weight_buffer = std::unique_ptr( + ggml_backend_dev_buffer_from_host_ptr(ggml_backend_get_device(hrx_backend), + second_weight.data(), second_weight.size(), + second_weight.size()), ggml_backend_buffer_free); + REQUIRE(hrx_second_weight_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(hrx_second_weight_buffer.get(), hrx_second_weight, + second_weight.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(hrx_second_weight_buffer.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); @@ -4781,13 +7149,10 @@ static void run_vector_q8_publish_cpu_reference_case() { REQUIRE(hrx_buffer != nullptr); const std::vector first_weight = make_matmul_weight_bytes(GGML_TYPE_F32, input_size, intermediate_size, 6); - const std::vector second_weight = - make_matmul_weight_bytes(GGML_TYPE_Q4_K, intermediate_size, output_size, 11); const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); set_tensor_pair_bytes(cpu_backend, cpu_first_weight, hrx_backend, hrx_first_weight, first_weight.data(), first_weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_second_weight, hrx_backend, hrx_second_weight, second_weight.data(), - second_weight.size()); + ggml_backend_tensor_set(cpu_second_weight, second_weight.data(), 0, second_weight.size()); set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); @@ -4880,7 +7245,8 @@ static void run_dense_matmul_postops_cpu_reference_case(ggml_type weight_type if (include_next_rmsnorm) { require_kernel_subsequence( scheduled_kernel_sequence(hrx_graph), - { expected_kernel, expected_plain_rmsnorm_binary_kernel(output_size, token_count), "loom_libs:ggml_binary_f32" }); + { expected_kernel, expected_plain_rmsnorm_binary_kernel(output_size, token_count), + "loom_libs:ggml_binary_f32" }); } else { require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); } @@ -6383,15 +8749,15 @@ static void register_dense_matmul_cases(Suite & suite) { for (const ggml_type type : { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_IQ2_S, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, GGML_TYPE_F16, GGML_TYPE_BF16 }) { - const char * vector_kernel = type == GGML_TYPE_IQ2_S ? "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32" : - type == GGML_TYPE_Q6_K ? "loom_libs:ggml_mul_mat_vector_q6_f32_f32" : - "loom_libs:ggml_mul_mat_vector_f32_f32"; + const char * vector_kernel = type == GGML_TYPE_IQ2_S ? "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32" : + "loom_libs:ggml_mul_mat_vector_f32_f32"; const int64_t outputs = type == GGML_TYPE_IQ2_S ? 640 : (type == GGML_TYPE_IQ4_NL ? 1024 : 128); const int64_t input = type == GGML_TYPE_IQ4_NL ? 640 : kQwenHiddenSize; - const char * route = type == GGML_TYPE_IQ4_NL ? "tiled" : "skinny"; - const char * kernel = type == GGML_TYPE_IQ2_S ? "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32" : - type == GGML_TYPE_IQ4_NL ? "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" : - "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; + const bool tiled_route = type == GGML_TYPE_IQ2_S || type == GGML_TYPE_IQ4_NL; + const char * route = tiled_route ? "tiled" : "skinny"; + const char * kernel = type == GGML_TYPE_IQ2_S ? "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32" : + type == GGML_TYPE_IQ4_NL ? "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; suite.device_case( "dense_matmul." + type_name(type) + "." + route + ".tokens2.outputs" + std::to_string(outputs), [type, outputs, input, kernel] { run_dense_matmul_cpu_reference_case(type, kernel, 2, outputs, input); }); @@ -6481,7 +8847,7 @@ static void register_dense_matmul_cases(Suite & suite) { [op, outputs] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, outputs, op); + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 4, outputs, op); }); } for (const int64_t tokens : { 256, 512 }) { @@ -6491,7 +8857,7 @@ static void register_dense_matmul_cases(Suite & suite) { [op] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 512, 4096, op); + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 4096, op); }); } else { for (const int64_t outputs : { 128, 4096 }) { @@ -6500,7 +8866,7 @@ static void register_dense_matmul_cases(Suite & suite) { [op, outputs] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, outputs, op); }); } @@ -6511,7 +8877,7 @@ static void register_dense_matmul_cases(Suite & suite) { suite.device_case( "dense_matmul_swiglu.q4_k.q4_k.lowtoken.tokens" + std::to_string(tokens) + ".outputs4160", [tokens] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", tokens, 4160); }); } @@ -6520,30 +8886,50 @@ static void register_dense_matmul_cases(Suite & suite) { "dense_matmul_swiglu.q4_k.q4_k.swiglu.tokens" + std::to_string(tokens) + ".outputs16384.input4096", [tokens] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); }); } suite.device_case("dense_matmul_swiglu.q4_k.q4_k.reglu.tokens3.outputs16384.input4352", [] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 3, 16384, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 3, 16384, GGML_GLU_OP_REGLU, 4352); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.lowtoken.tokens4.outputs128", [] { run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128); + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 4, 128); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs128", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 128); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 128); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs4096", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 4096); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens1024.outputs2048", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 1024, 2048); + }); + suite.device_case("imported_packed_gated_q4.lowtoken", [] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + 4, 128, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, true); + }); + suite.device_case("imported_packed_gated_q4.prefill", [] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", + 256, 128, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, true); + }); + suite.device_case("imported_packed_binary_q4.lowtoken", [] { + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + 4, 128, true); + }); + suite.device_case("imported_packed_binary_q4.prefill", [] { + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", + 256, 128, true); }); suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, @@ -6654,11 +9040,11 @@ static void register_dense_matmul_cases(Suite & suite) { }); } suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs256.input5120", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 256, 5120, + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 5120, true); }); suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs320.input512", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 320, 512, + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 320, 512, true); }); } @@ -6741,6 +9127,47 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("rmsnorm_mul.hidden256.tokens4", [] { run_rmsnorm_mul_case(256, 4); }); suite.device_case("rmsnorm_mul.hidden2048.tokens1", [] { run_rmsnorm_mul_case(2048, 1); }); suite.device_case("router_projection.tokens4", [] { run_router_projection_case(4); }); + suite.device_case("imported_native_weight", [] { run_imported_native_weight_case(); }); + suite.device_case("mutable_host_input_after_writeback.cpu", [] { run_mutable_host_input_after_writeback_case(false); }); + suite.device_case("mutable_host_input_after_writeback.hrx", [] { run_mutable_host_input_after_writeback_case(true); }); + suite.host_case("weight_demand_comparison", [] { run_weight_demand_comparison_case(); }); + suite.device_case("weight_reclamation_sequence", [] { run_weight_reclamation_sequence_case(); }); + suite.device_case("weight_reclamation_overlap", [] { run_weight_reclamation_overlap_case(); }); + suite.device_case("failed_submission_weight_lifetime", [] { run_failed_submission_weight_lifetime_case(); }); + suite.device_case("imported_packed_weight.q4_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q4_K); }); + suite.device_case("imported_packed_weight.q6_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q6_K); }); + suite.device_case("mixed_weight_layout_program.q4_k", [] { run_mixed_weight_layout_program_case(GGML_TYPE_Q4_K); }); + suite.device_case("mixed_weight_layout_program.q6_k", [] { run_mixed_weight_layout_program_case(GGML_TYPE_Q6_K); }); + // Standard GGUF weights are immutable: native host copies can be retained across token inputs. + suite.device_case("cpu_host_native_weight.q4_k", [] { + run_imported_packed_weight_case(GGML_TYPE_Q4_K, PackedWeightSourceKind::CpuHost); + }); + suite.device_case("cpu_host_native_weight.q6_k", [] { + run_imported_packed_weight_case(GGML_TYPE_Q6_K, PackedWeightSourceKind::CpuHost); + }); + suite.device_case("hrx_host_native_weight.q4_k", [] { + run_imported_packed_weight_case(GGML_TYPE_Q4_K, PackedWeightSourceKind::HrxHost); + }); + suite.device_case("coherent_subregion_native_weight.q4_k", [] { + run_imported_packed_weight_case(GGML_TYPE_Q4_K, PackedWeightSourceKind::CoherentSubregion); + }); + suite.device_case("native_cpu_weight_import_transition.q4_k", [] { + run_native_cpu_weight_during_import_layout_transition_case(GGML_TYPE_Q4_K); + }); + suite.device_case("native_cpu_weight_import_transition.q6_k", [] { + run_native_cpu_weight_during_import_layout_transition_case(GGML_TYPE_Q6_K); + }); + suite.device_case("q4_prefill_source_eligibility", [] { run_q4_prefill_source_eligibility_case(); }); + suite.device_case("uploaded_native_q6_vector", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128, + 2048, false, false, GGML_TYPE_COUNT, true); + }); + suite.device_case("uploaded_native_gated_q4", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, + GGML_GLU_OP_SWIGLU, 2048, true); + }); + suite.device_case("mixed_symmetric_source_eligibility", [] { run_mixed_symmetric_source_eligibility_case(); }); suite.device_case("router_top8.tokens4", [] { run_router_top8_case(4); }); for (int64_t experts : { 128, 256, 512 }) { for (int64_t tokens : { 1, 2, 512 }) { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index aff6a7c90b68..055dc950c079 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -7,6 +7,10 @@ #include "llama.h" #include "llama-cpp.h" #include "../src/llama-ext.h" +#ifdef GGML_USE_HRX +#include "ggml-hrx.h" +#include "../src/llama-model.h" +#endif #include #include @@ -156,6 +160,7 @@ static void usage(char ** argv) { LOG(" -o, --out Save generated test models to instead of running backend tests\n"); LOG(" -v Set log verbosity level\n"); LOG(" -b, --backend Run only on the given backend device\n"); + LOG(" --hrx-import Test mmap-backed HRX weight import\n"); LOG(" -h, --help Show this help message\n\n"); LOG("Examples:\n"); LOG(" %s\n", argv[0]); @@ -226,7 +231,7 @@ struct kv_writer { } }; -static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { +static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool hrx_compatible = false) { gguf_context_ptr ret(gguf_init_empty()); kv_writer ms(arch, ret.get()); const uint32_t n_ctx = 256; @@ -236,6 +241,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { uint32_t n_head = 2; uint32_t n_ff = 384; uint32_t n_layer = 2; + if (hrx_compatible) { + n_ff = 256; + } if (arch == LLM_ARCH_LLAMA4) { n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4 } else if (arch == LLM_ARCH_MAION_CODER) { @@ -1368,6 +1376,97 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con return all_ok ? 0 : 1; } +static int test_hrx_mmap_import(const size_t seed, const float stdev) { +#ifndef GGML_USE_HRX + GGML_UNUSED(seed); + GGML_UNUSED(stdev); + throw std::runtime_error("mmap import test requires an HRX build"); +#else + ggml_backend_load_all(); + ggml_backend_dev_t hrx_device = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t device = ggml_backend_dev_get(i); + if (std::strncmp(ggml_backend_dev_name(device), "HRX", 3) == 0) { + hrx_device = device; + break; + } + } + if (hrx_device == nullptr) { + throw std::runtime_error("HRX device unavailable for mmap import test"); + } + ggml_backend_dev_props props = {}; + ggml_backend_dev_get_props(hrx_device, &props); + if (!props.caps.buffer_from_host_ptr) { + throw std::runtime_error("HRX host-pointer import is not advertised"); + } + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_LLAMA, false, true); + auto cpu_model = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {}); + const ggml_tensor * cpu_weight = cpu_model.first->get_tensor("blk.0.attn_q.weight"); + if (cpu_weight == nullptr) { + throw std::runtime_error("tiny GGUF has no attention weight"); + } + std::vector expected(ggml_nbytes(cpu_weight)); + ggml_backend_tensor_get(cpu_weight, expected.data(), 0, expected.size()); + + const std::filesystem::path path = std::filesystem::temp_directory_path() / + ("hrx-import-" + std::to_string(seed) + ".gguf"); + struct RemoveFile { + std::filesystem::path path; + ~RemoveFile() { std::error_code ignored; std::filesystem::remove(path, ignored); } + } remove_file{ path }; + llama_model_save_to_file(cpu_model.first.get(), path.string().c_str()); + + ggml_backend_dev_t devices[] = { hrx_device, nullptr }; + llama_model_params model_params = llama_model_default_params(); + model_params.devices = devices; + model_params.n_gpu_layers = 99; + model_params.load_mode = LLAMA_LOAD_MODE_MMAP; + llama_model_ptr imported(llama_model_load_from_file(path.string().c_str(), model_params)); + if (!imported) { + throw std::runtime_error("failed to load mmap-backed HRX model"); + } + // The loader retains the mapped bytes needed by imported buffers after the file is unlinked. + if (!std::filesystem::remove(path)) { + throw std::runtime_error("failed to unlink mmap-backed GGUF before readback"); + } + llama_context_params context_params = llama_context_default_params(); + context_params.n_ctx = 0; + context_params.n_ubatch = 64; + llama_context_ptr imported_context(llama_init_from_model(imported.get(), context_params)); + if (!imported_context) { + throw std::runtime_error("failed to initialize mmap-backed HRX context"); + } + const ggml_tensor * imported_weight = imported->get_tensor("blk.0.attn_q.weight"); + if (imported_weight == nullptr) { + throw std::runtime_error("mmap-backed HRX attention weight is missing"); + } + if (imported_weight->buffer == nullptr) { + throw std::runtime_error("mmap-backed HRX attention weight is unbound"); + } + ggml_backend_dev_t weight_device = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(imported_weight->buffer)); + if (weight_device != hrx_device) { + throw std::runtime_error("tiny attention weight was not placed on HRX"); + } + if (ggml_backend_buffer_is_host(imported_weight->buffer)) { + throw std::runtime_error("mmap-backed HRX attention weight was classified as host storage"); + } + if (!ggml_backend_hrx_buffer_is_imported(imported_weight->buffer)) { + throw std::runtime_error("tiny attention weight was not imported from mmap"); + } + std::vector actual(expected.size()); + ggml_backend_tensor_get(imported_weight, actual.data(), 0, actual.size()); + if (actual != expected) { + throw std::runtime_error("mmap-backed HRX weight differs after loader cleanup"); + } + const std::vector logits = get_logits(imported.get(), imported_context.get(), {1}); + if (logits.empty()) { + throw std::runtime_error("mmap-backed HRX model produced no logits"); + } + return 0; +#endif +} + int main(int argc, char ** argv) { // init the logger at max verbosity. filter with a custom callback respecting the user-configure verbosity common_log_set_verbosity_thold(LOG_LEVEL_DEBUG); @@ -1381,6 +1480,7 @@ int main(int argc, char ** argv) { float stdev = 0.1f; std::string out; const char * target_backend = nullptr; + bool hrx_import = false; int verbosity = LOG_LEVEL_ERROR; @@ -1448,6 +1548,8 @@ int main(int argc, char ** argv) { usage(argv); return 1; } + } else if (strcmp(argv[i], "--hrx-import") == 0) { + hrx_import = true; } else { LOG_ERR("%s: unknown argument: %s\n", __func__, argv[i]); usage(argv); @@ -1463,6 +1565,9 @@ int main(int argc, char ** argv) { init_normal_pool(seed, stdev); try { + if (hrx_import) { + return test_hrx_mmap_import(seed, stdev); + } if (!out.empty()) { return save_models(arch_filter, seed, stdev, verbosity, out); } From 65c906a34cc6dcfb00c67a49e16a2b44bbd2ec45 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Sun, 1 Nov 2026 12:59:07 +0000 Subject: [PATCH 02/18] ggml-hrx: enable mmap for automatic model loading Assisted-by: Codex --- ggml/src/ggml-hrx/ggml-hrx.cpp | 1 + tests/test-llama-archs.cpp | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index cbc4a15be8d8..34fe022cf023 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -653,6 +653,7 @@ static void device_props(ggml_backend_dev_t device, ggml_backend_dev_props * pro props->type = GGML_BACKEND_DEVICE_TYPE_GPU; props->device_id = nullptr; props->caps = { true, true, true, false }; + props->caps.mmap_support = true; } static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parameters) { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 055dc950c079..f2f4a05e6e12 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -1421,7 +1421,10 @@ static int test_hrx_mmap_import(const size_t seed, const float stdev) { llama_model_params model_params = llama_model_default_params(); model_params.devices = devices; model_params.n_gpu_layers = 99; - model_params.load_mode = LLAMA_LOAD_MODE_MMAP; + // Keep AUTO so advertised device capabilities must enable mmap import. + if (model_params.load_mode != LLAMA_LOAD_MODE_AUTO) { + throw std::runtime_error("mmap import test requires the default AUTO load mode"); + } llama_model_ptr imported(llama_model_load_from_file(path.string().c_str(), model_params)); if (!imported) { throw std::runtime_error("failed to load mmap-backed HRX model"); From 9ac50b8435678f65dc51556ab4271d46c6ae9492 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Sun, 1 Nov 2026 22:25:14 +0000 Subject: [PATCH 03/18] ggml-hrx: trust admitted warm weight bindings Reuse exact prepared owners at the admitted residency revision while retaining asynchronous completion ownership. Share immutable weight lease lists across prepared, recorded and current owners, and preserve direct host-input completion. Bead: lemonade-ws-tek.26 --- .../runtime/command-program-executor.cpp | 19 +- .../runtime/command-program-executor.h | 6 +- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 103 +++------ .../ggml-hrx/runtime/graph-program-cache.cpp | 211 +++++------------- .../ggml-hrx/runtime/graph-program-cache.h | 26 +-- ggml/src/ggml-hrx/runtime/graph-replay.h | 10 + .../prepared-command-program-cache.cpp | 15 +- .../runtime/weight-submission-tracker.h | 11 +- tests/hrx-backend-test.cpp | 65 +++++- tests/hrx-buffer-test.cpp | 3 +- tests/hrx-ops-test.cpp | 187 ++++++++++++++-- 11 files changed, 377 insertions(+), 279 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 06991016e68e..78a806459367 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -788,6 +788,9 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec const CommandProgram & commands, const CommandProgramBindings & bindings, PreparedCommandProgram & prepared) { + // Publish the owner before acquisition so partial preparation retains every acquired lease. + auto resident_weights = std::make_shared>(); + prepared.resident_host_weights = resident_weights; std::vector materialized; Status status; materialized.reserve(bindings.bindings().size()); @@ -834,7 +837,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec device_binding.offset = 0; device_binding.length = resident.lease.length(); device_binding.capacity = resident.lease.length(); - prepared.resident_host_weights.push_back(std::move(resident.lease)); + resident_weights->push_back(std::move(resident.lease)); } materialized.push_back(std::move(device_binding)); } @@ -872,7 +875,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec } const bool can_report_residency = status.success() && context.host_weights != nullptr; if (can_report_residency) { - context.host_weights->report_coexisting_representations(prepared.resident_host_weights, + context.host_weights->report_coexisting_representations(*prepared.resident_host_weights, "same command program"); } return CommandProgramBindings::from_bindings(std::move(materialized), status); @@ -2264,6 +2267,18 @@ bool execute_prepared_command_program(const CommandProgramExecutionContext & con !execute_prepared_command_list_serial(context, commands.commands, "main", commands.commands.size())) { return false; } + // Registered host inputs can be overwritten as soon as graph_compute returns; finish their queued reads first. + Status flush_status = flush_stream_commands(context, "flush direct HRX command program"); + if (!flush_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(flush_status)); + return false; + } + Status wait_status = wait_stream_commands(context, "wait for direct HRX command program"); + if (!wait_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(wait_status)); + return false; + } + mark_graph_replay_synchronized(context); HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-download-pre"); Status download_status = download_prepared_host_staging(context, commands); if (!download_status.success()) { diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index 211953fa5b4a..e2306f191d98 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -136,7 +136,8 @@ struct PreparedCommandProgram { std::vector initialization_commands; std::vector commands; std::vector host_staging; - std::vector resident_host_weights; + std::shared_ptr> resident_host_weights = + std::make_shared>(); std::vector program_constants; Status status; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; @@ -148,7 +149,8 @@ struct RecordedCommandGraph { hrx_graph_t graph = nullptr; hrx_graph_exec_t exec = nullptr; std::vector retained_buffers; - std::vector resident_host_weights; + std::shared_ptr> resident_host_weights = + std::make_shared>(); std::vector retained_executables; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; size_t dispatch_count = 0; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 4bbd4b3eaadf..18b5a94d2cbf 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -9,7 +9,6 @@ #include "runtime/transient-arena.h" #include -#include #include #include @@ -134,6 +133,7 @@ void WeightSubmissionTracker::publish_demand(const std::vector for (const auto & [region, layouts] : retained_regions_) { current_demand_.insert(current_demand_.end(), layouts.begin(), layouts.end()); } + ++revision_; } Status WeightSubmissionTracker::poll_completed(bool * completed) { @@ -198,7 +198,7 @@ Status WeightSubmissionTracker::record(hrx_device_t device, hrx_stream_t stream, if (demand_changed) { publish_demand(demand); } - current_leases_.clear(); + release_current_leases(); } return status; } @@ -275,7 +275,7 @@ std::vector WeightSubmissionTracker::outstanding_leases() const if (pending.resources == nullptr) { continue; } - const auto & resident = pending.resources->prepared.resident_host_weights; + const auto & resident = *pending.resources->prepared.resident_host_weights; leases.insert(leases.end(), resident.begin(), resident.end()); } return leases; @@ -362,16 +362,21 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { return result; } - const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); + const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); + const bool admitted_weights = use_graph_prepared && + lookup.program->admitted_at(context_.weight_submissions.revision()); + const bool reusable_owner = use_graph_prepared && admitted_weights; + const bool mutable_sources_current = reusable_owner && lookup.program->mutable_sources_current(); + const bool trusted_weights = reusable_owner && mutable_sources_current; + if (use_graph_prepared) { + lookup.program->invalidate_admission(); + } GraphProgramMatch binding_match = std::move(lookup.match); const bool has_prepared = lookup.program->has_prepared_program(); - const bool may_use_partial_bindings = use_graph_prepared && has_prepared; - bool full_bindings = true; - if (may_use_partial_bindings) { - const bool imported_sources_current = lookup.program->imported_sources_current(); - binding_match = imported_sources_current ? lookup.program->match_host_staging_graph(graph) : - lookup.program->match_trusted_graph(graph, true); - full_bindings = !imported_sources_current; + const bool needs_prepared_match = use_graph_prepared && has_prepared; + if (needs_prepared_match) { + binding_match = trusted_weights ? lookup.program->match_host_staging_graph(graph) : + lookup.program->match_trusted_graph(graph, true); } if (!binding_match.valid()) { result.status.append(binding_match.status); @@ -383,51 +388,19 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { result.status.append(bindings.status); return result; } - std::shared_ptr> cached_demand; - if (!full_bindings) { - const GraphProgramMatch candidate_match = lookup.program->match_weight_demand_graph(graph); - if (!candidate_match.valid()) { - result.status.append(candidate_match.status); - return result; - } - CommandProgramBindings candidate_bindings = bind_external_value_buffers(candidate_match); - if (!candidate_bindings.valid()) { - result.status.append(candidate_bindings.status); - return result; - } - cached_demand = lookup.program->cached_weight_demand_for(candidate_bindings); - } - std::optional full_demand_bindings; - const CommandProgramBindings * demand_bindings = &bindings; - const bool needs_full_demand_bindings = !full_bindings && cached_demand == nullptr; - if (needs_full_demand_bindings) { - const GraphProgramMatch demand_match = lookup.program->match_trusted_graph(graph, true); - if (!demand_match.valid()) { - result.status.append(demand_match.status); - return result; - } - full_demand_bindings.emplace(bind_external_value_buffers(demand_match)); - if (!full_demand_bindings->valid()) { - result.status.append(full_demand_bindings->status); - return result; - } - demand_bindings = &*full_demand_bindings; - } - std::vector collected_demand; - if (cached_demand == nullptr) { + std::vector current_demand; + if (!trusted_weights) { const CommandProgramWeightDemands weight_demands = - bind_command_program_weight_demands(lookup.program->weight_requests(), *demand_bindings); + bind_command_program_weight_demands(lookup.program->weight_requests(), bindings); if (!weight_demands.valid()) { result.status.append(weight_demands.status); return result; } for (const CommandProgramWeightDemand & demand : weight_demands.weights) { - collected_demand.push_back(demand.source); + current_demand.push_back(demand.source); } } - const std::vector & current_demand = - cached_demand != nullptr ? *cached_demand : collected_demand; - const bool demand_changed = !context_.weight_submissions.demand_matches(current_demand); + const bool demand_changed = !trusted_weights && !context_.weight_submissions.demand_matches(current_demand); std::vector merged_demand; if (demand_changed) { merged_demand = context_.weight_submissions.merged_demand_for(current_demand); @@ -442,26 +415,6 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { context_.graph_programs.invalidate_obsolete_weights(merged_demand); context_.prepared_programs.invalidate_obsolete_weights(merged_demand); } - const bool prepared_program_missing = !lookup.program->has_prepared_program(); - const bool restore_full_bindings = use_graph_prepared && prepared_program_missing; - if (restore_full_bindings) { - if (full_demand_bindings) { - bindings = std::move(*full_demand_bindings); - full_bindings = true; - } else if (!full_bindings) { - const GraphProgramMatch full_match = lookup.program->match_trusted_graph(graph, true); - if (!full_match.valid()) { - result.status.append(full_match.status); - return result; - } - bindings = bind_external_value_buffers(full_match); - if (!bindings.valid()) { - result.status.append(bindings.status); - return result; - } - full_bindings = true; - } - } const bool reconcile_before = demand_changed || completed_before; if (reconcile_before) { const auto & retained = demand_changed ? merged_demand : context_.weight_submissions.current_demand(); @@ -480,7 +433,7 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { &context_.graph_replay_state, }; const PreparedCommandProgramCacheExecutionResult execution = - use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings, full_bindings) : + use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings, !trusted_weights) : context_.prepared_programs.execute_with_result(execution_context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), bindings); @@ -490,12 +443,9 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { demand_changed); result.status.append(completion_status); const bool admitted = execution.success && completion_status.success(); - const bool cache_needs_refresh = cached_demand == nullptr || restore_full_bindings; - const bool cache_can_refresh = admitted && use_graph_prepared && cache_needs_refresh; - if (cache_can_refresh) { - const CommandProgramBindings & cache_bindings = - full_bindings ? bindings : *full_demand_bindings; - lookup.program->remember_weight_demand(cache_bindings, current_demand); + const bool publish_admission = admitted && use_graph_prepared; + if (publish_admission) { + lookup.program->admit(execution.execution_resources, context_.weight_submissions.revision(), trusted_weights); } if (!admitted) { const auto & retained_demand = context_.weight_submissions.current_demand(); @@ -528,6 +478,9 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { lookup.program->mark_imported_writes(graph); } if (!result.status.success()) { + if (use_graph_prepared) { + lookup.program->invalidate_admission(); + } return result; } if (!execution.success) { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index f63d4255561f..9ff1b101ffd3 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -20,31 +20,6 @@ namespace ggml::hrx { namespace { -static bool same_weight_demand_binding(const CommandProgramBinding & left, const CommandProgramBinding & right) { - const bool same_value = left.value == right.value; - const bool same_weight_class = left.weight == right.weight; - const bool same_binding_class = same_value && same_weight_class; - if (!same_binding_class) { - return false; - } - if (!left.weight) { - // Mutable graph inputs can change each token without changing which weights are demanded. - return true; - } - const bool same_storage = left.buffer == right.buffer && left.host_data == right.host_data && - left.offset == right.offset && left.length == right.length && - left.capacity == right.capacity && left.empty_value == right.empty_value; - const bool same_source = left.identity == right.identity && left.generation == right.generation && - left.source_host_data == right.source_host_data && - left.source_offset == right.source_offset && left.source_capacity == right.source_capacity && - left.source_type == right.source_type && left.source_ne == right.source_ne && - left.source_nb == right.source_nb; - const bool same_layout = left.layout == right.layout && left.layout_source_type == right.layout_source_type && - left.input_size == right.input_size && left.output_size == right.output_size && - left.source_length == right.source_length; - return same_storage && same_source && same_layout; -} - static const ggml_tensor * tensor_storage_root(const ggml_tensor * tensor) { while (tensor != nullptr && tensor->view_src != nullptr) { tensor = tensor->view_src; @@ -430,16 +405,6 @@ GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_ } } if (!bind_external) { - for (const GraphProgramExternalSlot & slot : external_slots_) { - const ggml_tensor * tensor = resolve_external_slot(current_graph, slot, result.status); - if (tensor == nullptr) { - return result; - } - if (ggml_backend_hrx_weight_layout_eligible(tensor) != slot.weight_layout_eligible) { - result.status.log("external weight source class changed for value %d", slot.value.value); - return result; - } - } return result; } result.external_bindings.reserve(external_slots_.size()); @@ -480,39 +445,16 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur return result; } -GraphProgramMatch GraphProgram::match_weight_demand_graph(const ggml_cgraph & current_graph) const { - GraphProgramMatch result; - result.external_bindings.reserve(weight_demand_slots_.size()); - for (const GraphProgramExternalSlot & slot : weight_demand_slots_) { - const ggml_tensor * tensor = resolve_external_slot(current_graph, slot, result.status); - if (tensor == nullptr) { - return result; - } - const bool source_class_current = - ggml_backend_hrx_weight_layout_eligible(tensor) == slot.weight_layout_eligible; - if (!source_class_current) { - result.status.log("external weight source class changed for value %d", slot.value.value); - return result; - } - result.external_bindings.push_back({ slot.value, tensor }); - } - return result; -} - Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { Status status; external_slots_.clear(); - weight_demand_slots_.clear(); - weight_demand_values_.clear(); - weight_demand_bindings_.clear(); - weight_demand_.reset(); - weight_demand_resources_.reset(); - imported_source_watches_.clear(); - written_values_.clear(); + admitted_resources_.reset(); + mutable_source_watches_.clear(); + written_slots_.clear(); + std::unordered_set written_values; external_slot_by_value_.clear(); fast_path_nodes_ = graph.nodes; external_slots_.reserve(match.external_bindings.size()); - std::unordered_set read_values; auto collect_access = [&](const std::vector & commands) { for (const Command & command : commands) { for (const CommandBinding & binding : command.bindings) { @@ -522,12 +464,7 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra const bool writes_value = writes_only || reads_and_writes; const bool graph_write = graph_value && writes_value; if (graph_write) { - written_values_.insert(binding.value.value); - } - const bool reads_value = binding.access != ResourceAccess::Write; - const bool graph_read = graph_value && reads_value; - if (graph_read) { - read_values.insert(binding.value.value); + written_values.insert(binding.value.value); } } } @@ -538,13 +475,12 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra ValueBufferBinding value_binding; const bool bound = ggml_backend_hrx_resolve_value_buffer(binding.tensor, value_binding); const bool weight = bound && value_binding.weight; - const bool host_source = bound && value_binding.source_host_data != nullptr; - const bool tracks_generation = weight && host_source; ggml_backend_hrx_imported_region region = {}; - if (ggml_backend_hrx_get_retained_region(binding.tensor, region)) { - // Watch each physical region, including aliases in one buffer, but only weight contents. - imported_source_watches_.push_back({ binding.tensor, region.identity, region.generation, - region.offset, region.length, tracks_generation }); + const bool mutable_source = bound && !weight; + const bool retained_source = mutable_source && ggml_backend_hrx_get_retained_region(binding.tensor, region); + if (retained_source) { + // Mutable inputs may rebind under one graph UID; immutable weight bindings are trusted. + mutable_source_watches_.push_back({ binding.tensor, region.identity, region.offset, region.length }); } GraphProgramExternalSlot slot; slot.value = binding.value; @@ -580,14 +516,8 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra } external_slot_by_value_[binding.value.value] = external_slots_.size(); external_slots_.push_back(slot); - } - for (const GraphProgramExternalSlot & slot : external_slots_) { - const bool read_by_command = read_values.count(slot.value.value) != 0; - const bool written_by_command = written_values_.count(slot.value.value) != 0; - const bool candidate = read_by_command && !written_by_command; - if (candidate) { - weight_demand_slots_.push_back(slot); - weight_demand_values_.insert(slot.value.value); + if (written_values.count(slot.value.value) != 0) { + written_slots_.push_back(slot); } } return status; @@ -600,12 +530,46 @@ bool GraphProgram::has_prepared_program() const { bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { // The UID prevents graph-arena address reuse from passing the fast path. - return graph.uid == uid_ && graph.nodes == fast_path_nodes_; + const bool allocated_graph = graph.uid != 0; + const bool same_uid = graph.uid == uid_; + const bool same_nodes = graph.nodes == fast_path_nodes_; + return allocated_graph && same_uid && same_nodes; +} + +bool GraphProgram::admitted_at(uint64_t revision) const { + std::lock_guard lock(prepared_mutex_); + const bool prepared = has_prepared_; + const bool same_owner = admitted_resources_.lock() == resources_; + const bool current_revision = admitted_revision_ == revision; + return prepared && same_owner && current_revision; } -bool GraphProgram::imported_sources_current() const { +void GraphProgram::invalidate_admission() { std::lock_guard lock(prepared_mutex_); - for (const ImportedSourceWatch & watch : imported_source_watches_) { + admitted_resources_.reset(); +} + +void GraphProgram::admit(const std::shared_ptr & resources, + uint64_t revision, bool reused_demand) { + std::lock_guard lock(prepared_mutex_); + const bool prepared = has_prepared_; + const bool same_owner = resources == resources_; + const bool current_owner = prepared && same_owner; + if (!current_owner) { + return; + } + admitted_resources_ = resources; + admitted_revision_ = revision; + if (reused_demand) { + ++weight_demand_hits_; + } else { + ++weight_demand_builds_; + } +} + +bool GraphProgram::mutable_sources_current() const { + std::lock_guard lock(prepared_mutex_); + for (const MutableSourceWatch & watch : mutable_source_watches_) { ggml_backend_hrx_imported_region region = {}; if (!ggml_backend_hrx_get_retained_region(watch.tensor, region)) { return false; @@ -613,12 +577,10 @@ bool GraphProgram::imported_sources_current() const { const bool same_identity = region.identity == watch.identity; const bool same_offset = region.offset == watch.offset; const bool same_length = region.length == watch.length; - const bool same_generation = !watch.tracks_content_generation || region.generation == watch.generation; const bool identity_changed = !same_identity; const bool offset_changed = !same_offset; const bool length_changed = !same_length; - const bool generation_changed = !same_generation; - const bool source_changed = identity_changed || offset_changed || length_changed || generation_changed; + const bool source_changed = identity_changed || offset_changed || length_changed; if (source_changed) { return false; } @@ -628,11 +590,7 @@ bool GraphProgram::imported_sources_current() const { void GraphProgram::mark_imported_writes(const ggml_cgraph & graph) const { std::unordered_set marked_identities; - for (const GraphProgramExternalSlot & slot : external_slots_) { - const bool written_value = written_values_.count(slot.value.value) != 0; - if (!written_value) { - continue; - } + for (const GraphProgramExternalSlot & slot : written_slots_) { Status status; const ggml_tensor * tensor = resolve_external_slot(graph, slot, status); if (tensor == nullptr) { @@ -654,56 +612,6 @@ PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { return prepared_stats_; } -std::shared_ptr> GraphProgram::cached_weight_demand_for( - const CommandProgramBindings & bindings) const { - std::lock_guard lock(prepared_mutex_); - const bool prepared_current = has_prepared_ && resources_ != nullptr; - const bool cached_for_prepared = weight_demand_resources_.lock() == resources_; - const bool has_cached_demand = weight_demand_ != nullptr; - const bool reusable_demand = prepared_current && cached_for_prepared && has_cached_demand; - if (!reusable_demand) { - return nullptr; - } - size_t index = 0; - for (const CommandProgramBinding & binding : bindings.bindings()) { - if (weight_demand_values_.count(binding.value.value) == 0) { - continue; - } - const bool has_cached_binding = index < weight_demand_bindings_.size(); - const bool same_binding = has_cached_binding && - same_weight_demand_binding(binding, weight_demand_bindings_[index]); - if (!same_binding) { - return nullptr; - } - ++index; - } - if (index != weight_demand_bindings_.size()) { - return nullptr; - } - ++weight_demand_hits_; - return weight_demand_; -} - -void GraphProgram::remember_weight_demand(const CommandProgramBindings & bindings, - const std::vector & demand) { - std::vector candidates; - candidates.reserve(weight_demand_slots_.size()); - for (const CommandProgramBinding & binding : bindings.bindings()) { - if (weight_demand_values_.count(binding.value.value) != 0) { - candidates.push_back(binding); - } - } - std::lock_guard lock(prepared_mutex_); - const bool prepared_current = has_prepared_ && resources_ != nullptr; - if (!prepared_current) { - return; - } - weight_demand_bindings_ = std::move(candidates); - weight_demand_ = std::make_shared>(demand); - weight_demand_resources_ = resources_; - ++weight_demand_builds_; -} - std::pair GraphProgram::weight_demand_stats() const { std::lock_guard lock(prepared_mutex_); return { weight_demand_builds_, weight_demand_hits_ }; @@ -715,13 +623,14 @@ void GraphProgram::invalidate_obsolete_weights(const std::vectorprepared.resident_host_weights) { + for (const HostWeightLease & lease : *resources_->prepared.resident_host_weights) { const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), [&](const HostWeightSource & source) { return lease.matches(source); }); obsolete |= !demanded; } if (obsolete) { resources_.reset(); + admitted_resources_.reset(); has_prepared_ = false; prepared_bindings_fingerprint_.clear(); } @@ -732,6 +641,8 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( const CommandProgramBindings & bindings, bool full_bindings) { PreparedCommandProgramCacheExecutionResult result; + std::lock_guard lock(prepared_mutex_); + admitted_resources_.reset(); if (commands_ == nullptr || !commands_->valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; result.graph_replay_ineligible_reason = "invalid_graph_program"; @@ -742,7 +653,6 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( return result; } - std::lock_guard lock(prepared_mutex_); const std::string bindings_fingerprint = full_bindings ? command_program_bindings_fingerprint(bindings).value : std::string{}; const bool fingerprint_changed = prepared_bindings_fingerprint_ != bindings_fingerprint; @@ -757,6 +667,7 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } context.graph_replay_state->mark_stream_synchronized(); resources_.reset(); + admitted_resources_.reset(); has_prepared_ = false; } if (!has_prepared_) { @@ -769,11 +680,10 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } has_prepared_ = true; prepared_bindings_fingerprint_ = bindings_fingerprint; - for (ImportedSourceWatch & watch : imported_source_watches_) { + for (MutableSourceWatch & watch : mutable_source_watches_) { ggml_backend_hrx_imported_region region = {}; if (ggml_backend_hrx_get_retained_region(watch.tensor, region)) { watch.identity = region.identity; - watch.generation = region.generation; watch.offset = region.offset; watch.length = region.length; } @@ -784,9 +694,12 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } result.execution_resources = resources_; - if (environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_REPLAY") || debug_serial_command_execution_enabled()) { + const bool replay_disabled = graph_replay_disabled_by_environment(); + const bool serial_execution = debug_serial_command_execution_enabled(); + const bool execute_directly = replay_disabled || serial_execution; + if (execute_directly) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = debug_serial_command_execution_enabled() ? + result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : "disabled_by_environment"; result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, resources_->prepared); diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index 94e85f2cc642..e12f2f2064ec 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -73,17 +73,16 @@ class GraphProgram { GraphProgramMatch match_current_graph(const ggml_cgraph & graph) const; GraphProgramMatch match_trusted_graph(const ggml_cgraph & graph, bool bind_external = true) const; GraphProgramMatch match_host_staging_graph(const ggml_cgraph & graph) const; - GraphProgramMatch match_weight_demand_graph(const ggml_cgraph & graph) const; - std::shared_ptr> cached_weight_demand_for( - const CommandProgramBindings & bindings) const; - void remember_weight_demand(const CommandProgramBindings & bindings, - const std::vector & demand); + bool admitted_at(uint64_t revision) const; + void invalidate_admission(); + void admit(const std::shared_ptr & resources, + uint64_t revision, bool reused_demand); Status capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match); bool has_prepared_program() const; bool can_use_prepared_fast_path(const ggml_cgraph & graph) const; - bool imported_sources_current() const; + bool mutable_sources_current() const; void mark_imported_writes(const ggml_cgraph & graph) const; PreparedCommandProgramCacheStats prepared_stats() const; @@ -110,27 +109,22 @@ class GraphProgram { uint64_t command_shape_hash_ = 0; std::vector external_slots_; - std::vector weight_demand_slots_; - std::unordered_set weight_demand_values_; - struct ImportedSourceWatch { + struct MutableSourceWatch { const ggml_tensor * tensor; uint64_t identity; - uint64_t generation; size_t offset; size_t length; - bool tracks_content_generation; }; - std::vector imported_source_watches_; - std::unordered_set written_values_; + std::vector mutable_source_watches_; + std::vector written_slots_; std::unordered_map external_slot_by_value_; const ggml_tensor * const * fast_path_nodes_ = nullptr; mutable std::mutex prepared_mutex_; std::shared_ptr resources_; std::string prepared_bindings_fingerprint_; PreparedCommandProgramCacheStats prepared_stats_; - std::vector weight_demand_bindings_; - std::shared_ptr> weight_demand_; - std::weak_ptr weight_demand_resources_; + std::weak_ptr admitted_resources_; + uint64_t admitted_revision_ = 0; mutable uint64_t weight_demand_builds_ = 0; mutable uint64_t weight_demand_hits_ = 0; bool has_prepared_ = false; diff --git a/ggml/src/ggml-hrx/runtime/graph-replay.h b/ggml/src/ggml-hrx/runtime/graph-replay.h index 3bcd3abdcbbe..5a9cd4ed4144 100644 --- a/ggml/src/ggml-hrx/runtime/graph-replay.h +++ b/ggml/src/ggml-hrx/runtime/graph-replay.h @@ -2,6 +2,8 @@ #include #include +#include +#include namespace ggml::hrx { @@ -35,6 +37,14 @@ inline const char * hrx_graph_replay_event_name(HrxGraphReplayEvent event) { return "unknown"; } +inline bool graph_replay_disabled_by_environment() { + const char * value = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); + const bool set = value != nullptr; + const bool nonempty = set && value[0] != '\0'; + const bool enabled = nonempty && std::strcmp(value, "0") != 0; + return enabled; +} + inline uint64_t hrx_graph_replay_now_ns() { return static_cast( std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index 2968100f6d33..c4edd68e08ca 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -127,13 +127,17 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ } std::lock_guard entry_lock(entry->mutex); + const bool replay_disabled = graph_replay_disabled_by_environment(); + const bool serial_execution = debug_serial_command_execution_enabled(); + const bool execute_directly = replay_disabled || serial_execution; if (entry->resources != nullptr && entry->resources->prepared.valid()) { auto & resources = *entry->resources; result.execution_resources = entry->resources; record_hit(); - if (debug_serial_command_execution_enabled()) { + if (execute_directly) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = "debug_serial_execution"; + result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : + "disabled_by_environment"; result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources.prepared); if (!result.success) { result.status.log("execute cached HRX command program failed"); @@ -171,9 +175,10 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ result.execution_resources = resources; record_build(); - if (debug_serial_command_execution_enabled()) { + if (execute_directly) { result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = "debug_serial_execution"; + result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : + "disabled_by_environment"; result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); if (!result.success) { result.status.log("execute prepared HRX command program failed"); @@ -227,7 +232,7 @@ void PreparedCommandProgramCache::invalidate_obsolete_weights(const std::vector< continue; } bool obsolete = false; - for (const HostWeightLease & lease : entry->resources->prepared.resident_host_weights) { + for (const HostWeightLease & lease : *entry->resources->prepared.resident_host_weights) { const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), [&](const HostWeightSource & source) { return lease.matches(source); }); obsolete |= !demanded; diff --git a/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h b/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h index 0573b2def49e..1bbcd10ea39b 100644 --- a/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h +++ b/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h @@ -32,10 +32,14 @@ class WeightSubmissionTracker { void fail_completion_record_for_test() { fail_next_record_ = true; } const std::vector & current_demand() const { return current_demand_; } + uint64_t revision() const { return revision_; } bool demand_matches(const std::vector & demand) const; std::vector merged_demand_for(const std::vector & demand) const; - void release_current_leases() { current_leases_.clear(); } - const std::vector & current_leases() const { return current_leases_; } + void release_current_leases() { current_leases_.reset(); } + const std::vector & current_leases() const { + static const std::vector empty; + return current_leases_ != nullptr ? *current_leases_ : empty; + } std::vector outstanding_leases() const; private: @@ -69,7 +73,8 @@ class WeightSubmissionTracker { std::vector current_demand_; RegionDemands retained_regions_; - std::vector current_leases_; + std::shared_ptr> current_leases_; + uint64_t revision_ = 0; std::vector pending_; bool fail_next_record_ = false; }; diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 655b12ec17de..be3f75cd68ff 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -11777,7 +11777,7 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { return staging.value == unused_value->id.value; })); REQUIRE(prepared.host_staging.size() == 2); - REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(prepared.resident_host_weights->size() == 1); REQUIRE(state->host_weights.stats().resident_bytes == weight_bytes); REQUIRE(state->host_weights.stats().misses == 1); // The vector matmul declares both host input and output as reads. @@ -11799,7 +11799,7 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { auto prepared = ggml::hrx::prepare_command_program(execution_context, initialization_only, bindings); REQUIRE(prepared.valid()); REQUIRE(prepared.host_staging.size() == 2); - REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(prepared.resident_host_weights->size() == 1); REQUIRE(state->host_weights.stats().misses == 1); } ggml_backend_buffer_free(cpu_buffer); @@ -14649,7 +14649,68 @@ static void run_initializer_host_input_ownership() { ggml_free(ctx); } +static void run_registered_host_input_reuse_waits_before_mutation() { + constexpr int64_t count = 128; + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * state = static_cast(backend->context); + ggml_context * ctx = ggml_init({ 256 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, count); + ggml_tensor * output_a = ggml_add(ctx, input, input); + ggml_tensor * output_b = ggml_add(ctx, input, input); + REQUIRE(output_a != nullptr); + REQUIRE(output_b != nullptr); + ggml_cgraph * graph_a = ggml_new_graph(ctx); + ggml_cgraph * graph_b = ggml_new_graph(ctx); + REQUIRE(graph_a != nullptr); + REQUIRE(graph_b != nullptr); + ggml_build_forward_expand(graph_a, output_a); + ggml_build_forward_expand(graph_b, output_b); + graph_a->uid = 84027; + graph_b->uid = 84028; + + const bool direct_host_binding = state->device->use_direct_host_bindings; + state->device->use_direct_host_bindings = false; + ggml_backend_buffer_type_t host_type = ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend)); + ggml_backend_buffer_t host = ggml_backend_alloc_ctx_tensors_from_buft(ctx, host_type); + state->device->use_direct_host_bindings = direct_host_binding; + REQUIRE(host != nullptr); + ggml::hrx::ValueBufferBinding resolved_input; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(input, resolved_input)); + REQUIRE(resolved_input.host_data != nullptr); + REQUIRE(resolved_input.buffer == nullptr); + + ggml::hrx::GraphExecutor executor(*state); + auto * input_data = static_cast(input->data); + auto * first_output = static_cast(output_a->data); + auto * second_output = static_cast(output_b->data); + std::fill_n(input_data, count, 0.5f); + REQUIRE(executor.execute(*graph_a).success()); + REQUIRE(executor.execute(*graph_b).success()); + ggml_backend_synchronize(backend); + + std::fill_n(first_output, count, -1000.0f); + std::fill_n(second_output, count, -1000.0f); + std::fill_n(input_data, count, 1.0f); + REQUIRE(executor.execute(*graph_a).success()); + std::fill_n(input_data, count, 3.0f); + REQUIRE(executor.execute(*graph_b).success()); + ggml_backend_synchronize(backend); + for (int64_t i = 0; i < count; ++i) { + REQUIRE(first_output[i] == 2.0f); + REQUIRE(second_output[i] == 6.0f); + } + + ggml_backend_buffer_free(host); + ggml_free(ctx); + ggml_backend_free(backend); +} + static void register_hrx_backend_device_cases(test_runner::Suite & suite) { + suite.device_case("registered_host_input_reuse_waits_before_mutation", [] { + run_registered_host_input_reuse_waits_before_mutation(); + }); suite.device_case("initializer_host_input_ownership", [] { run_initializer_host_input_ownership(); }); suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index a4f782753394..f8b3afd9fde9 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -592,7 +592,8 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { ggml::hrx::RecordedCommandGraph recorded; hrx_buffer_retain(first.lease.buffer()); recorded.retained_buffers.push_back(first.lease.buffer()); - recorded.resident_host_weights.push_back(first.lease); + recorded.resident_host_weights = + std::make_shared>(std::initializer_list{ first.lease }); weights.clear(); weight_stats = weights.stats(); REQUIRE(weight_stats.hits == 1); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 128194363ab3..e13cb6063d2c 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3205,6 +3205,7 @@ static void run_weight_reclamation_sequence_case() { }; const auto native_first = execute(hrx_native, hrx_rows, expected_native); REQUIRE(native_first.weight_materializations == 1); + REQUIRE(state->weight_submissions.revision() == 1); const auto packed_first = execute(hrx_packed, hrx_matmul, expected_packed); REQUIRE(packed_first.weight_materializations == 2); REQUIRE(packed_first.weight_peak_bytes == @@ -3214,7 +3215,9 @@ static void run_weight_reclamation_sequence_case() { REQUIRE(native_second.weight_materializations == 3); const auto rebuilt_native_demand = state->graph_programs.stats(); REQUIRE(rebuilt_native_demand.weight_demand_builds > before_native_rebuild.weight_demand_builds); + const uint64_t warm_revision = state->weight_submissions.revision(); const auto native_unchanged = execute(hrx_native, hrx_rows, expected_native); + REQUIRE(state->weight_submissions.revision() == warm_revision); REQUIRE(native_unchanged.weight_materializations == native_second.weight_materializations); REQUIRE(native_unchanged.weight_uploaded_bytes == native_second.weight_uploaded_bytes); REQUIRE(native_unchanged.prepared_program_hits > native_second.prepared_program_hits); @@ -3335,6 +3338,48 @@ static void run_weight_reclamation_sequence_case() { completion_fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows); completion_fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second); + // Logical demand can survive a discarded prepared owner and its physical allocation. + execute(hrx_packed, hrx_matmul, expected_packed, 3); + const auto packed_lookup = state->graph_programs.get_or_build( + *hrx_packed, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(packed_lookup.valid()); + const uint64_t retained_revision = state->weight_submissions.revision(); + REQUIRE(packed_lookup.program->admitted_at(retained_revision)); + const auto before_owner_reset = state->host_weights.stats(); + const auto retained_demand = state->weight_submissions.current_demand(); + ggml_backend_synchronize(backend); + state->graph_programs.invalidate_obsolete_weights({}); + state->prepared_programs.invalidate_obsolete_weights({}); + state->weight_submissions.release_current_leases(); + state->host_weights.evict_except({}, {}); + REQUIRE(!packed_lookup.program->has_prepared_program()); + REQUIRE(!packed_lookup.program->admitted_at(retained_revision)); + REQUIRE(state->host_weights.stats().allocation_count == 0); + REQUIRE(state->weight_submissions.demand_matches(retained_demand)); + REQUIRE(state->weight_submissions.revision() == retained_revision); + execute(hrx_packed, hrx_matmul, expected_packed, 1); + REQUIRE(state->host_weights.stats().misses == before_owner_reset.misses + 1); + REQUIRE(state->weight_submissions.revision() == retained_revision); + REQUIRE(packed_lookup.program->admitted_at(retained_revision)); + + // A failed capture of an otherwise warm execution cannot leave its admission reusable. + const auto before_capture_failure = state->graph_programs.stats(); + const auto current_leases = state->weight_submissions.current_leases(); + state->weight_submissions.fail_completion_record_for_test(); + REQUIRE(ggml_backend_graph_compute(backend, hrx_packed) == GGML_STATUS_FAILED); + REQUIRE(!packed_lookup.program->admitted_at(retained_revision)); + REQUIRE(state->weight_submissions.revision() == retained_revision); + REQUIRE(state->graph_programs.stats().weight_demand_builds == before_capture_failure.weight_demand_builds); + REQUIRE(state->graph_programs.stats().weight_demand_hits == before_capture_failure.weight_demand_hits); + REQUIRE(state->weight_submissions.current_leases().size() == current_leases.size()); + for (size_t i = 0; i < current_leases.size(); ++i) { + REQUIRE(state->weight_submissions.current_leases()[i].buffer() == current_leases[i].buffer()); + } + execute(hrx_packed, hrx_matmul, expected_packed, 1); + REQUIRE(state->weight_submissions.revision() == retained_revision); + REQUIRE(packed_lookup.program->admitted_at(retained_revision)); + REQUIRE(state->graph_programs.stats().weight_demand_builds == before_capture_failure.weight_demand_builds + 1); + ggml_backend_buffer_free(hrx_tensors); ggml_backend_buffer_free(cpu_tensors); ggml_backend_buffer_free(imported); @@ -3351,7 +3396,7 @@ static void run_weight_reclamation_overlap_case() { const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); const bool had_flag = old_flag != nullptr; const std::string old_flag_value = had_flag ? old_flag : ""; - REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "1", 1) == 0); + REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "0", 1) == 0); ggml_backend_t cpu = init_cpu_backend(); ggml_backend_t backend = ggml_backend_hrx_init(0); @@ -3478,6 +3523,82 @@ static void run_weight_reclamation_overlap_case() { ggml_backend_hrx_cache_stats cross_before = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cross_before)); REQUIRE(cross_before.weight_duplicate_warnings == 0); + // The public executors finish their kernel reads before returning. Queue warmed recordings directly + // to exercise ownership of genuinely outstanding reads without weakening that completion boundary. + struct Submission { + std::shared_ptr resources; + std::vector demand; + }; + auto prepare_submission = [&](ggml_cgraph * graph) { + auto lookup = state->graph_programs.get_or_build( + *graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(lookup.valid()); + auto match = lookup.program->match_trusted_graph(*graph, true); + REQUIRE(match.valid()); + std::vector external_bindings; + for (const auto & external : match.external_bindings) { + ggml::hrx::ValueBufferBinding buffer; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); + std::array ne = {}; + std::array nb = {}; + for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { + ne[dimension] = external.tensor->ne[dimension]; + nb[dimension] = external.tensor->nb[dimension]; + } + external_bindings.push_back(ggml::hrx::make_command_program_binding( + external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); + auto demands = ggml::hrx::collect_command_program_weight_demands(lookup.program->commands(), bindings); + REQUIRE(demands.valid()); + const ggml::hrx::CommandProgramExecutionContext context = { + state->device->device, state->stream, state->device->architecture.c_str(), + &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, + &state->host_transfers, &state->host_weights, &state->device->host_buffers, + &state->graph_replay_state, + }; + auto execution = lookup.program->execute_with_result(context, bindings, true); + REQUIRE(execution.success); + REQUIRE(execution.execution_resources->recorded.valid()); + REQUIRE(execution.execution_resources->recorded.resident_host_weights == + execution.execution_resources->prepared.resident_host_weights); + Submission submission; + submission.resources = std::move(execution.execution_resources); + for (const auto & demand : demands.weights) { + submission.demand.push_back(demand.source); + } + return submission; + }; + auto queue_submission = [&](Submission & submission) { + const bool changed = !state->weight_submissions.demand_matches(submission.demand); + if (changed) { + state->weight_submissions.release_current_leases(); + } + const auto launch = ggml::hrx::take_status(hrx_graph_exec_launch( + submission.resources->recorded.exec, state->stream)); + if (launch) { + return false; + } + const auto * weight_list = submission.resources->prepared.resident_host_weights.get(); + const auto capture = state->weight_submissions.record( + state->device->device, state->stream, submission.resources, submission.demand, true, changed); + const bool shares_weights = &state->weight_submissions.current_leases() == weight_list; + submission.resources.reset(); + if (!capture.success()) { + return false; + } + const auto & retained = state->weight_submissions.current_demand(); + state->graph_programs.invalidate_obsolete_weights(retained); + state->prepared_programs.invalidate_obsolete_weights(retained); + const auto outstanding = state->weight_submissions.outstanding_leases(); + state->host_weights.evict_except(retained, outstanding); + state->host_weights.report_coexisting_representations(outstanding, "test outstanding submissions"); + return shares_weights; + }; + auto cross_native_submission = prepare_submission(hrx_native_graph); + auto cross_packed_submission = prepare_submission(hrx_packed_graph); + const std::vector poisoned_packed(expected_packed.size(), 123.0f); + ggml_backend_tensor_set(hrx_packed, poisoned_packed.data(), 0, poisoned_packed.size() * sizeof(float)); hrx_semaphore_t cross_gate = nullptr; REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &cross_gate))); REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { cross_gate, 1 }))); @@ -3493,9 +3614,9 @@ static void run_weight_reclamation_overlap_case() { hrx_status_ignore(signal); } }); - const bool cross_native = ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + const bool cross_native = queue_submission(cross_native_submission); const bool cross_packed = cross_native && - ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_SUCCESS; + queue_submission(cross_packed_submission); ggml_backend_hrx_cache_stats cross_queued = {}; const bool cross_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cross_queued); bool cross_obsolete_completed = true; @@ -3530,6 +3651,7 @@ static void run_weight_reclamation_overlap_case() { // A cold preparation failure must leave an already queued native reader's lease intact. REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); + auto cold_native_submission = prepare_submission(hrx_native_graph); const std::vector poisoned_native(expected_native.size(), 123.0f); ggml_backend_tensor_set(hrx_rows, poisoned_native.data(), 0, poisoned_native.size() * sizeof(float)); @@ -3549,7 +3671,7 @@ static void run_weight_reclamation_overlap_case() { } }); const bool native_before_failure = - ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + queue_submission(cold_native_submission); state->host_weights.fail_cold_acquire_after_for_test(0); const bool cold_failed = ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_FAILED; @@ -3572,6 +3694,12 @@ static void run_weight_reclamation_overlap_case() { REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); REQUIRE(warmed.weight_live_allocations == 2); + auto mixed_submission = prepare_submission(hrx_mixed_graph); + auto native_submission = prepare_submission(hrx_native_graph); + auto packed_submission = prepare_submission(hrx_packed_graph); + const std::vector poisoned_mixed(expected_mixed.size(), 123.0f); + ggml_backend_tensor_set(hrx_mixed, poisoned_mixed.data(), 0, poisoned_mixed.size() * sizeof(float)); + ggml_backend_tensor_set(hrx_packed, poisoned_packed.data(), 0, poisoned_packed.size() * sizeof(float)); hrx_semaphore_t gate = nullptr; REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); @@ -3588,11 +3716,11 @@ static void run_weight_reclamation_overlap_case() { hrx_status_ignore(signal); } }); - const bool mixed_queued = ggml_backend_graph_compute_async(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS; + const bool mixed_queued = queue_submission(mixed_submission); const bool native_queued = mixed_queued && - ggml_backend_graph_compute_async(backend, hrx_native_graph) == GGML_STATUS_SUCCESS; + queue_submission(native_submission); const bool packed_queued = native_queued && - ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_SUCCESS; + queue_submission(packed_submission); bool stream_complete = true; const bool stream_query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); ggml_backend_hrx_cache_stats queued = {}; @@ -3674,7 +3802,7 @@ static void run_failed_submission_weight_lifetime_case() { const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); const bool had_flag = old_flag != nullptr; const std::string old_flag_value = had_flag ? old_flag : ""; - REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "1", 1) == 0); + REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "0", 1) == 0); ggml_backend_t cpu = init_cpu_backend(); ggml_backend_t backend = ggml_backend_hrx_init(0); @@ -3775,6 +3903,18 @@ static void run_failed_submission_weight_lifetime_case() { REQUIRE(preuploaded.valid()); ggml_backend_synchronize(backend); preuploaded.lease = {}; + const ggml::hrx::CommandProgramExecutionContext execution_context = { + state->device->device, state->stream, state->device->architecture.c_str(), + &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, + &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, + }; + auto execution = packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); + REQUIRE(execution.success); + REQUIRE(execution.execution_resources->recorded.valid()); + ggml_backend_synchronize(backend); + ggml_backend_tensor_set(hrx_output, poisoned.data(), 0, poisoned.size() * sizeof(float)); + // The invalid final command fails after the real weight reader is queued, before the completion boundary. + execution.execution_resources->prepared.commands.emplace_back(); hrx_semaphore_t gate = nullptr; REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); @@ -3790,13 +3930,8 @@ static void run_failed_submission_weight_lifetime_case() { hrx_status_ignore(signal); } }); - const ggml::hrx::CommandProgramExecutionContext execution_context = { - state->device->device, state->stream, state->device->architecture.c_str(), - &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, - &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, - }; - auto execution = packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); - const bool command_queued = execution.success && execution.execution_resources != nullptr; + const bool execution_failed = !ggml::hrx::execute_prepared_command_program( + execution_context, execution.execution_resources->prepared); std::weak_ptr sole_resource = execution.execution_resources; const bool has_execution_resources = execution.execution_resources != nullptr; const bool record_ok = has_execution_resources && tracker.record(state->device->device, state->stream, @@ -3819,7 +3954,7 @@ static void run_failed_submission_weight_lifetime_case() { gate_cv.notify_one(); const bool signal_ok = !ggml::hrx::take_status(hrx_semaphore_signal(gate, 1)); watchdog.join(); - REQUIRE(command_queued && record_ok && poll_ok && query_ok && signal_ok && !timed_out); + REQUIRE(execution_failed && record_ok && poll_ok && query_ok && signal_ok && !timed_out); REQUIRE(!completed && !stream_complete && demand_unchanged && tracker_only_owner && resource_retained); REQUIRE(pending_stats.allocation_count == 1); REQUIRE(pending_stats.resident_bytes == packed_source.materialized_length); @@ -4060,7 +4195,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, auto prepared = ggml::hrx::prepare_command_program( execution_context, lookup.program->commands(), bindings); REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights.size() == expected_materializations); + REQUIRE(prepared.resident_host_weights->size() == expected_materializations); const bool staged_weight = std::any_of(prepared.host_staging.begin(), prepared.host_staging.end(), [&](const auto & staging) { return staging.value == weight_value->id.value; }); REQUIRE(!staged_weight); @@ -4133,7 +4268,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, auto lookup = backend_state->graph_programs.get_or_build( *hrx_decode, ggml::hrx::get_qwen_kernel_corpus(), backend_state->device->architecture); REQUIRE(lookup.valid()); - REQUIRE(lookup.program->imported_sources_current()); + REQUIRE(lookup.program->mutable_sources_current()); auto bindings_for_match = [](const ggml::hrx::GraphProgramMatch & match) { std::vector external; for (const auto & slot : match.external_bindings) { @@ -4214,7 +4349,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, auto prepared = ggml::hrx::prepare_command_program( execution_context, lookup.program->commands(), value_map_bindings); REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(prepared.resident_host_weights->size() == 1); ggml::hrx::RecordedCommandGraph recorded; const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( execution_context, lookup.program->commands(), value_map_bindings, prepared, recorded); @@ -4258,7 +4393,9 @@ static void run_imported_packed_weight_case(ggml_type weight_type, replay_context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), full_bindings); REQUIRE(cache_replay.success); - REQUIRE(cache_replay.graph_replay_event == ggml::hrx::HrxGraphReplayEvent::Hit); + REQUIRE(cache_replay.graph_replay_event == + (ggml::hrx::graph_replay_disabled_by_environment() ? ggml::hrx::HrxGraphReplayEvent::Disabled : + ggml::hrx::HrxGraphReplayEvent::Hit)); ggml_backend_synchronize(hrx_backend); } std::vector readback(source.size()); @@ -4287,7 +4424,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, const std::vector shifted_input = make_pattern_f32(kInputSize, 79, 0.06f); set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, shifted_input.data(), shifted_input.size() * sizeof(float)); - REQUIRE(!lookup.program->imported_sources_current()); + REQUIRE(!lookup.program->mutable_sources_current()); compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); ggml_backend_hrx_cache_stats after_offset_change = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &after_offset_change)); @@ -4302,7 +4439,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, const std::vector rebound_input = make_pattern_f32(kInputSize, 83, 0.06f); set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, rebound_input.data(), rebound_input.size() * sizeof(float)); - REQUIRE(!lookup.program->imported_sources_current()); + REQUIRE(!lookup.program->mutable_sources_current()); ggml_backend_hrx_cache_stats before_rebind = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_rebind)); compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); @@ -4349,7 +4486,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, }; auto prepared = ggml::hrx::prepare_command_program(execution_context, commands, bindings); REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights.size() == 1); + REQUIRE(prepared.resident_host_weights->size() == 1); ggml::hrx::RecordedCommandGraph recorded; const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( execution_context, commands, bindings, prepared, recorded); @@ -4947,11 +5084,13 @@ static void run_q4_prefill_source_eligibility_case() { REQUIRE(structural_replay.graph_program_builds == cache.graph_program_builds); REQUIRE(structural_replay.prepared_program_hits > cache.prepared_program_hits); - // Keep the graph and node pointers fixed while changing only the live weight binding. + // Source rebinding starts a new graph identity; bindings within an allocated graph stay stable. Fixture & replay = fixtures[0]; ggml_backend_buffer_t imported_buffer = replay.weight->buffer; void * imported_data = replay.weight->data; + uint64_t rebound_uid = 960010; for (const bool imported : { true, false, true }) { + replay.graph->uid = rebound_uid++; replay.weight->buffer = imported ? imported_buffer : fixtures[1].weight->buffer; replay.weight->data = imported ? imported_data : fixtures[1].weight->data; ggml_backend_hrx_cache_stats before = {}; From b7b75b89d60b5008c9503a71154e63797209a9b7 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 00:35:48 +0000 Subject: [PATCH 04/18] hrx: retain first-use weight materializations for backend lifetime Assisted-by: Codex --- ggml/src/ggml-hrx/README.md | 4 +- ggml/src/ggml-hrx/backend-context.h | 2 - ggml/src/ggml-hrx/ggml-hrx.cpp | 17 - .../runtime/command-program-executor.cpp | 5 - .../runtime/command-program-executor.h | 3 +- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 355 +------------ .../ggml-hrx/runtime/graph-program-cache.cpp | 69 --- .../ggml-hrx/runtime/graph-program-cache.h | 18 - ggml/src/ggml-hrx/runtime/host-memory.cpp | 156 ++---- ggml/src/ggml-hrx/runtime/host-memory.h | 10 +- .../prepared-command-program-cache.cpp | 27 +- .../runtime/prepared-command-program-cache.h | 1 - .../runtime/weight-submission-tracker.h | 82 --- tests/hrx-buffer-test.cpp | 26 +- tests/hrx-ops-test.cpp | 492 ++++++------------ 15 files changed, 231 insertions(+), 1036 deletions(-) delete mode 100644 ggml/src/ggml-hrx/runtime/weight-submission-tracker.h diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index cfa97b1331d9..3e98c3b60dfe 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -8,7 +8,9 @@ Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx Live weight bytes count physical allocations retained by the cache, a prepared program, or a recorded graph. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. -Each successfully admitted HRX submission updates the retained layout set only for the weight regions it reads. A later fragment using other weights, or no weights, leaves those regions resident. When a region is used with a different set, completed obsolete layouts are released while unfinished submissions keep their bindings until completion. A coherent native read also replaces that region's layout set, even when it needs no device allocation. Repeating warm disjoint fragments reuses their allocations; explicit cache clearing and backend teardown require materialization again. +Each immutable source region is materialized only when a command first requests its layout. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. A warm graph reuses its prepared bindings without rebuilding per-weight demand; mutable inputs still rebind and upload normally. + +Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. Keep borrowed weight bytes unchanged and alive while HRX can use them. A foreign CPU host buffer and its storage must stay alive for the lifetime of the HRX backend that caches it. Its identity depends on its wrapper and base addresses; destroying and recreating a buffer at those addresses while that backend lives can reuse old residency. Training and in-place weight updates are outside this path. Ordinary mutable inputs and outputs retain their host staging, writeback and synchronization behavior. diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index 1e9e211f2994..fdc8fbaf89e2 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -8,7 +8,6 @@ #include "runtime/kernel-executable-cache.h" #include "runtime/prepared-command-program-cache.h" #include "runtime/transient-arena.h" -#include "runtime/weight-submission-tracker.h" #include #include @@ -80,7 +79,6 @@ struct ggml_backend_hrx_context { ggml::hrx::TransientArena transient_arena; ggml::hrx::HostTransferManager host_transfers; ggml::hrx::HostWeightCache host_weights; - ggml::hrx::WeightSubmissionTracker weight_submissions; ggml::hrx::GraphReplayStreamState graph_replay_state; std::string name; }; diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 34fe022cf023..f448c1ba2194 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -584,14 +584,6 @@ static void backend_synchronize(ggml_backend_t backend) { auto * context = static_cast(backend->context); if (HRX_CHECK(hrx_stream_synchronize(context->stream))) { context->graph_replay_state.mark_stream_synchronized(); - const bool obsolete_completed = context->weight_submissions.has_obsolete_pending(); - context->weight_submissions.mark_stream_synchronized(); - if (obsolete_completed) { - const auto & current_demand = context->weight_submissions.current_demand(); - context->graph_programs.invalidate_obsolete_weights(current_demand); - context->prepared_programs.invalidate_obsolete_weights(current_demand); - context->host_weights.evict_except(current_demand, {}); - } } } @@ -1335,15 +1327,6 @@ bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_c return false; } auto * context = static_cast(backend->context); - bool completed = false; - const ggml::hrx::Status retirement = context->weight_submissions.poll_completed(&completed); - if (!retirement.success()) { - return false; - } - if (completed) { - context->host_weights.evict_except(context->weight_submissions.current_demand(), - context->weight_submissions.outstanding_leases()); - } const ggml::hrx::GraphProgramCacheStats graph_stats = context->graph_programs.stats(); const ggml::hrx::PreparedCommandProgramCacheStats prepared_stats = context->prepared_programs.stats(); const ggml::hrx::HostWeightCacheStats weight_stats = context->host_weights.stats(); diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 78a806459367..124f0f9183b0 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -873,11 +873,6 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec materialized.push_back(device_binding); prepared.host_staging.push_back(std::move(staging)); } - const bool can_report_residency = status.success() && context.host_weights != nullptr; - if (can_report_residency) { - context.host_weights->report_coexisting_representations(*prepared.resident_host_weights, - "same command program"); - } return CommandProgramBindings::from_bindings(std::move(materialized), status); } diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index e2306f191d98..c3091aa98409 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -168,8 +168,7 @@ struct RecordedCommandGraph { bool valid() const { return status.success() && exec != nullptr; } }; -// A submitted command may still use prepared bindings after its cache entry is retired. -// The execution completion tracker retains this object until the stream reaches the submission. +// Cached programs own their bindings across submissions; recorded handles are released before prepared leases. struct CachedCommandProgramResources { PreparedCommandProgram prepared; RecordedCommandGraph recorded; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 18b5a94d2cbf..c9895ae7644f 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -8,279 +8,11 @@ #include "runtime/prepared-command-program-cache.h" #include "runtime/transient-arena.h" -#include #include #include namespace ggml::hrx { -WeightSubmissionTracker::~WeightSubmissionTracker() { - mark_stream_synchronized(); -} - -bool WeightSubmissionTracker::RegionKey::operator==(const RegionKey & other) const { - return identity == other.identity && offset == other.offset && - length == other.length && type == other.type && ne == other.ne && nb == other.nb; -} - -size_t WeightSubmissionTracker::RegionKeyHash::operator()(const RegionKey & key) const { - uint64_t hash = UINT64_C(1469598103934665603); - auto mix = [&](uint64_t value) { - hash ^= value; - hash *= UINT64_C(1099511628211); - }; - mix(key.identity); - mix(key.offset); - mix(key.length); - mix(static_cast(key.type)); - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - mix(static_cast(key.ne[dimension])); - mix(key.nb[dimension]); - } - return static_cast(hash); -} - -WeightSubmissionTracker::RegionKey WeightSubmissionTracker::region_of(const HostWeightSource & source) { - return { source.identity, source.offset, source.region_length, source.region_type, - source.source_ne, source.source_nb }; -} - -static bool same_weight_representation(const HostWeightSource & left, const HostWeightSource & right) { - return left.generation == right.generation && left.capacity == right.capacity && - left.length == right.length && - left.materialized_length == right.materialized_length && left.layout == right.layout && - left.source_type == right.source_type && left.input_size == right.input_size && - left.output_size == right.output_size; -} - -WeightSubmissionTracker::RegionDemands WeightSubmissionTracker::group_demands( - const std::vector & demand) { - RegionDemands grouped; - for (const HostWeightSource & source : demand) { - auto & layouts = grouped[region_of(source)]; - const bool duplicate = std::any_of(layouts.begin(), layouts.end(), [&](const HostWeightSource & existing) { - return same_weight_representation(existing, source); - }); - if (!duplicate) { - layouts.push_back(source); - } - } - return grouped; -} - -bool WeightSubmissionTracker::demand_matches(const std::vector & demand) const { - bool needs_grouped_comparison = false; - for (const HostWeightSource & source : demand) { - const auto retained = retained_regions_.find(region_of(source)); - if (retained == retained_regions_.end()) { - return false; - } - const bool singleton_region = retained->second.size() == 1; - if (!singleton_region) { - needs_grouped_comparison = true; - break; - } - if (!same_weight_representation(retained->second.front(), source)) { - return false; - } - } - if (!needs_grouped_comparison) { - return true; - } - - const RegionDemands touched = group_demands(demand); - for (const auto & [region, layouts] : touched) { - const auto retained = retained_regions_.find(region); - const bool region_is_retained = retained != retained_regions_.end(); - if (!region_is_retained) { - return false; - } - const bool same_layout_count = layouts.size() == retained->second.size(); - if (!same_layout_count) { - return false; - } - for (const HostWeightSource & source : layouts) { - const bool present = std::any_of(retained->second.begin(), retained->second.end(), - [&](const HostWeightSource & existing) { - return same_weight_representation(existing, source); - }); - if (!present) { - return false; - } - } - } - return true; -} - -std::vector WeightSubmissionTracker::merged_demand_for( - const std::vector & demand) const { - RegionDemands merged = retained_regions_; - for (auto & [region, layouts] : group_demands(demand)) { - merged[region] = std::move(layouts); - } - std::vector result; - for (const auto & [region, layouts] : merged) { - result.insert(result.end(), layouts.begin(), layouts.end()); - } - return result; -} - -void WeightSubmissionTracker::publish_demand(const std::vector & demand) { - for (auto & [region, layouts] : group_demands(demand)) { - retained_regions_[region] = std::move(layouts); - } - current_demand_.clear(); - for (const auto & [region, layouts] : retained_regions_) { - current_demand_.insert(current_demand_.end(), layouts.begin(), layouts.end()); - } - ++revision_; -} - -Status WeightSubmissionTracker::poll_completed(bool * completed) { - Status status; - if (completed != nullptr) { - *completed = false; - } - for (auto pending = pending_.begin(); pending != pending_.end();) { - const bool has_event = pending->event != nullptr; - const bool has_timeline = pending->semaphore != nullptr; - const bool has_completion = has_event || has_timeline; - if (!has_completion) { - ++pending; - continue; - } - bool complete = false; - if (has_event) { - if (ErrorResult error = take_status(hrx_event_query(pending->event, &complete))) { - status.log("query HRX weight submission completion: %s", error->c_str()); - return status; - } - } else { - uint64_t completed_value = 0; - if (ErrorResult error = take_status(hrx_semaphore_query(pending->semaphore, &completed_value))) { - status.log("query HRX weight submission timeline: %s", error->c_str()); - return status; - } - complete = completed_value >= pending->value; - } - if (complete) { - const bool report_obsolete_completion = completed != nullptr && pending->may_hold_obsolete_demand; - if (report_obsolete_completion) { - *completed = true; - } - release_completion(*pending); - pending = pending_.erase(pending); - } else { - ++pending; - } - } - return status; -} - -void WeightSubmissionTracker::release_completion(Pending & pending) { - if (pending.event != nullptr) { - hrx_event_release(pending.event); - pending.event = nullptr; - } - if (pending.semaphore != nullptr) { - hrx_semaphore_release(pending.semaphore); - pending.semaphore = nullptr; - } -} - -Status WeightSubmissionTracker::record(hrx_device_t device, hrx_stream_t stream, - std::shared_ptr resources, - const std::vector & demand, bool admitted, - bool demand_changed) { - Status status; - if (resources == nullptr) { - if (admitted) { - if (demand_changed) { - publish_demand(demand); - } - release_current_leases(); - } - return status; - } - const bool admitted_new_demand = admitted && demand_changed; - if (admitted_new_demand) { - for (Pending & pending : pending_) { - pending.may_hold_obsolete_demand = true; - } - } - hrx_event_t event = nullptr; - hrx_semaphore_t semaphore = nullptr; - uint64_t value = 0; - if (fail_next_record_) { - fail_next_record_ = false; - status.log("injected HRX weight completion capture failure"); - } else if (admitted) { - // Capture after all queued work, including a possible host-download tail. - hrx_timeline_point_t position = {}; - if (ErrorResult error = take_status(hrx_stream_flush(stream))) { - status.log("flush HRX weight submission: %s", error->c_str()); - } else if (ErrorResult error = take_status(hrx_stream_get_timeline_position(stream, &position))) { - status.log("capture HRX weight submission timeline: %s", error->c_str()); - } else { - semaphore = position.semaphore; - value = position.value; - hrx_semaphore_retain(semaphore); - } - } else if (ErrorResult error = take_status(hrx_event_create(device, HRX_EVENT_FLAG_DISABLE_TIMING, &event))) { - status.log("create HRX weight submission completion: %s", error->c_str()); - } else if (ErrorResult error = take_status(hrx_event_record(event, stream))) { - status.log("record HRX weight submission completion: %s", error->c_str()); - } - if (!status.success()) { - if (event != nullptr) { - hrx_event_release(event); - } - if (semaphore != nullptr) { - hrx_semaphore_release(semaphore); - } - // A capture failure must not let queued commands outlive their bindings. - if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { - status.log("synchronize after HRX completion capture failure: %s", error->c_str()); - pending_.push_back({ nullptr, nullptr, 0, std::move(resources), true }); - return status; - } - mark_stream_synchronized(); - return status; - } - pending_.push_back({ event, semaphore, value, std::move(resources), !admitted }); - if (admitted) { - if (admitted_new_demand) { - publish_demand(demand); - } - current_leases_ = pending_.back().resources->prepared.resident_host_weights; - } - return status; -} - -void WeightSubmissionTracker::mark_stream_synchronized() { - for (Pending & pending : pending_) { - release_completion(pending); - } - pending_.clear(); -} - -bool WeightSubmissionTracker::has_obsolete_pending() const { - return std::any_of(pending_.begin(), pending_.end(), - [](const Pending & pending) { return pending.may_hold_obsolete_demand; }); -} - -std::vector WeightSubmissionTracker::outstanding_leases() const { - std::vector leases; - for (const Pending & pending : pending_) { - if (pending.resources == nullptr) { - continue; - } - const auto & resident = *pending.resources->prepared.resident_host_weights; - leases.insert(leases.end(), resident.begin(), resident.end()); - } - return leases; -} - GraphExecutor::GraphExecutor(ggml_backend_hrx_context & context) : context_(context) {} Status GraphExecutor::context_valid_for_graph_programs() const { @@ -363,18 +95,12 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { } const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); - const bool admitted_weights = use_graph_prepared && - lookup.program->admitted_at(context_.weight_submissions.revision()); - const bool reusable_owner = use_graph_prepared && admitted_weights; + const bool has_prepared = lookup.program->has_prepared_program(); + const bool reusable_owner = use_graph_prepared && has_prepared; const bool mutable_sources_current = reusable_owner && lookup.program->mutable_sources_current(); const bool trusted_weights = reusable_owner && mutable_sources_current; - if (use_graph_prepared) { - lookup.program->invalidate_admission(); - } GraphProgramMatch binding_match = std::move(lookup.match); - const bool has_prepared = lookup.program->has_prepared_program(); - const bool needs_prepared_match = use_graph_prepared && has_prepared; - if (needs_prepared_match) { + if (reusable_owner) { binding_match = trusted_weights ? lookup.program->match_host_staging_graph(graph) : lookup.program->match_trusted_graph(graph, true); } @@ -388,38 +114,6 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { result.status.append(bindings.status); return result; } - std::vector current_demand; - if (!trusted_weights) { - const CommandProgramWeightDemands weight_demands = - bind_command_program_weight_demands(lookup.program->weight_requests(), bindings); - if (!weight_demands.valid()) { - result.status.append(weight_demands.status); - return result; - } - for (const CommandProgramWeightDemand & demand : weight_demands.weights) { - current_demand.push_back(demand.source); - } - } - const bool demand_changed = !trusted_weights && !context_.weight_submissions.demand_matches(current_demand); - std::vector merged_demand; - if (demand_changed) { - merged_demand = context_.weight_submissions.merged_demand_for(current_demand); - } - bool completed_before = false; - result.status = context_.weight_submissions.poll_completed(&completed_before); - if (!result.status.success()) { - return result; - } - if (demand_changed) { - context_.weight_submissions.release_current_leases(); - context_.graph_programs.invalidate_obsolete_weights(merged_demand); - context_.prepared_programs.invalidate_obsolete_weights(merged_demand); - } - const bool reconcile_before = demand_changed || completed_before; - if (reconcile_before) { - const auto & retained = demand_changed ? merged_demand : context_.weight_submissions.current_demand(); - context_.host_weights.evict_except(retained, context_.weight_submissions.outstanding_leases()); - } const CommandProgramExecutionContext execution_context = { context_.device->device, context_.stream, @@ -438,51 +132,10 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { lookup.program->command_shape_hash(), lookup.program->commands(), bindings); result.graph_replay_event = execution.graph_replay_event; - const Status completion_status = context_.weight_submissions.record( - context_.device->device, context_.stream, execution.execution_resources, current_demand, execution.success, - demand_changed); - result.status.append(completion_status); - const bool admitted = execution.success && completion_status.success(); - const bool publish_admission = admitted && use_graph_prepared; - if (publish_admission) { - lookup.program->admit(execution.execution_resources, context_.weight_submissions.revision(), trusted_weights); - } - if (!admitted) { - const auto & retained_demand = context_.weight_submissions.current_demand(); - context_.graph_programs.invalidate_obsolete_weights(retained_demand); - context_.prepared_programs.invalidate_obsolete_weights(retained_demand); - } - bool completed_after = false; - const Status retirement_status = context_.weight_submissions.poll_completed(&completed_after); - result.status.append(retirement_status); - const bool reconcile_after = demand_changed || completed_after || !admitted; - if (reconcile_after) { - const std::vector outstanding = context_.weight_submissions.outstanding_leases(); - context_.host_weights.evict_except(context_.weight_submissions.current_demand(), outstanding); - const auto & current = context_.weight_submissions.current_leases(); - const bool outstanding_only_allocation = std::any_of(outstanding.begin(), outstanding.end(), - [&](const HostWeightLease & lease) { - return std::none_of(current.begin(), current.end(), - [&](const HostWeightLease & active) { return active.buffer() == lease.buffer(); }); - }); - const bool report_overlap = admitted && outstanding_only_allocation; - if (report_overlap) { - std::vector coexisting = current; - coexisting.insert(coexisting.end(), outstanding.begin(), outstanding.end()); - context_.host_weights.report_coexisting_representations(coexisting, - "current_vs_outstanding"); - } - } if (execution.execution_resources != nullptr) { - // A failed return can still follow queued writes or a successful fallback synchronization. + // A failed return can still follow queued writes. lookup.program->mark_imported_writes(graph); } - if (!result.status.success()) { - if (use_graph_prepared) { - lookup.program->invalidate_admission(); - } - return result; - } if (!execution.success) { result.status.append(execution.status); if (result.status.success()) { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index 9ff1b101ffd3..49fafd29de96 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -327,7 +327,6 @@ GraphProgram::GraphProgram(uint64_t uid, target_(std::move(target)), graph_(std::move(graph)), commands_(std::move(commands)), - weight_requests_(build_command_program_weight_requests(*commands_)), command_shape_(std::move(command_shape)), command_shape_hash_(command_program_shape_hash(command_shape_)) {} @@ -448,7 +447,6 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { Status status; external_slots_.clear(); - admitted_resources_.reset(); mutable_source_watches_.clear(); written_slots_.clear(); std::unordered_set written_values; @@ -536,37 +534,6 @@ bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { return allocated_graph && same_uid && same_nodes; } -bool GraphProgram::admitted_at(uint64_t revision) const { - std::lock_guard lock(prepared_mutex_); - const bool prepared = has_prepared_; - const bool same_owner = admitted_resources_.lock() == resources_; - const bool current_revision = admitted_revision_ == revision; - return prepared && same_owner && current_revision; -} - -void GraphProgram::invalidate_admission() { - std::lock_guard lock(prepared_mutex_); - admitted_resources_.reset(); -} - -void GraphProgram::admit(const std::shared_ptr & resources, - uint64_t revision, bool reused_demand) { - std::lock_guard lock(prepared_mutex_); - const bool prepared = has_prepared_; - const bool same_owner = resources == resources_; - const bool current_owner = prepared && same_owner; - if (!current_owner) { - return; - } - admitted_resources_ = resources; - admitted_revision_ = revision; - if (reused_demand) { - ++weight_demand_hits_; - } else { - ++weight_demand_builds_; - } -} - bool GraphProgram::mutable_sources_current() const { std::lock_guard lock(prepared_mutex_); for (const MutableSourceWatch & watch : mutable_source_watches_) { @@ -612,37 +579,12 @@ PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { return prepared_stats_; } -std::pair GraphProgram::weight_demand_stats() const { - std::lock_guard lock(prepared_mutex_); - return { weight_demand_builds_, weight_demand_hits_ }; -} - -void GraphProgram::invalidate_obsolete_weights(const std::vector & current_demand) { - std::lock_guard lock(prepared_mutex_); - if (!has_prepared_) { - return; - } - bool obsolete = false; - for (const HostWeightLease & lease : *resources_->prepared.resident_host_weights) { - const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), - [&](const HostWeightSource & source) { return lease.matches(source); }); - obsolete |= !demanded; - } - if (obsolete) { - resources_.reset(); - admitted_resources_.reset(); - has_prepared_ = false; - prepared_bindings_fingerprint_.clear(); - } -} - PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( const CommandProgramExecutionContext & context, const CommandProgramBindings & bindings, bool full_bindings) { PreparedCommandProgramCacheExecutionResult result; std::lock_guard lock(prepared_mutex_); - admitted_resources_.reset(); if (commands_ == nullptr || !commands_->valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; result.graph_replay_ineligible_reason = "invalid_graph_program"; @@ -667,7 +609,6 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( } context.graph_replay_state->mark_stream_synchronized(); resources_.reset(); - admitted_resources_.reset(); has_prepared_ = false; } if (!has_prepared_) { @@ -1010,9 +951,6 @@ GraphProgramCacheStats GraphProgramCache::stats() const { const PreparedCommandProgramCacheStats prepared = entry.second->prepared_stats(); stats.prepared_program_builds += prepared.builds; stats.prepared_program_hits += prepared.hits; - const auto [demand_builds, demand_hits] = entry.second->weight_demand_stats(); - stats.weight_demand_builds += demand_builds; - stats.weight_demand_hits += demand_hits; } return stats; } @@ -1024,13 +962,6 @@ void GraphProgramCache::clear() { last_program_ = nullptr; } -void GraphProgramCache::invalidate_obsolete_weights(const std::vector & current_demand) { - std::lock_guard lock(mutex_); - for (auto & program : programs_) { - program.second->invalidate_obsolete_weights(current_demand); - } -} - std::unique_ptr GraphProgramCache::build_program_from_imported(uint64_t uid, Graph && imported_graph, const KernelCorpus & corpus, diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index e12f2f2064ec..e67a348eaec2 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -68,16 +68,9 @@ class GraphProgram { const CommandProgram & commands() const { return *commands_; } - const CommandProgramWeightRequests & weight_requests() const { return weight_requests_; } - GraphProgramMatch match_current_graph(const ggml_cgraph & graph) const; GraphProgramMatch match_trusted_graph(const ggml_cgraph & graph, bool bind_external = true) const; GraphProgramMatch match_host_staging_graph(const ggml_cgraph & graph) const; - bool admitted_at(uint64_t revision) const; - void invalidate_admission(); - void admit(const std::shared_ptr & resources, - uint64_t revision, bool reused_demand); - Status capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match); bool has_prepared_program() const; @@ -86,9 +79,6 @@ class GraphProgram { void mark_imported_writes(const ggml_cgraph & graph) const; PreparedCommandProgramCacheStats prepared_stats() const; - std::pair weight_demand_stats() const; - - void invalidate_obsolete_weights(const std::vector & current_demand); PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, const CommandProgramBindings & bindings, @@ -104,7 +94,6 @@ class GraphProgram { std::string target_; std::unique_ptr graph_; std::unique_ptr commands_; - const CommandProgramWeightRequests weight_requests_; std::string command_shape_; uint64_t command_shape_hash_ = 0; @@ -123,10 +112,6 @@ class GraphProgram { std::shared_ptr resources_; std::string prepared_bindings_fingerprint_; PreparedCommandProgramCacheStats prepared_stats_; - std::weak_ptr admitted_resources_; - uint64_t admitted_revision_ = 0; - mutable uint64_t weight_demand_builds_ = 0; - mutable uint64_t weight_demand_hits_ = 0; bool has_prepared_ = false; }; @@ -135,8 +120,6 @@ struct GraphProgramCacheStats { uint64_t hits = 0; uint64_t prepared_program_builds = 0; uint64_t prepared_program_hits = 0; - uint64_t weight_demand_builds = 0; - uint64_t weight_demand_hits = 0; }; struct GraphProgramLookup { @@ -157,7 +140,6 @@ struct GraphProgramSupportResult { class GraphProgramCache { public: - void invalidate_obsolete_weights(const std::vector & current_demand); bool can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target) const; GraphProgramSupportResult check_support(const ggml_cgraph & graph, diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 091d443951a8..1900468b45cc 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -13,7 +13,6 @@ #include #include #include -#include #include #include #include @@ -1453,16 +1452,8 @@ struct HostWeightLease::Entry { hrx_buffer_t buffer = nullptr; size_t length = 0; std::string layout; - uint64_t source_identity = 0; - uint64_t source_generation = 0; - size_t source_capacity = 0; - size_t source_offset = 0; - size_t source_length = 0; - ggml_type source_type = GGML_TYPE_COUNT; - int64_t input_size = 0; - int64_t output_size = 0; - std::array source_ne = {}; - std::array source_nb = {}; + size_t region_length = 0; + ggml_type region_type = GGML_TYPE_COUNT; std::shared_ptr accounting; }; @@ -1485,20 +1476,6 @@ const std::string & HostWeightLease::layout() const { return entry_ != nullptr ? entry_->layout : empty; } -bool HostWeightLease::matches(const HostWeightSource & source) const { - if (!valid()) { - return false; - } - const Entry & entry = *entry_; - const size_t materialized_length = source.layout == kNativeWeightLayout ? source.length : source.materialized_length; - return entry.source_identity == source.identity && entry.source_generation == source.generation && - entry.source_capacity == source.capacity && entry.source_offset == source.offset && - entry.source_length == source.length && entry.length == materialized_length && - entry.layout == source.layout && entry.source_type == source.source_type && - entry.input_size == source.input_size && entry.output_size == source.output_size && - entry.source_ne == source.source_ne && entry.source_nb == source.source_nb; -} - HostWeightCache::HostWeightCache() : accounting_(std::make_shared()) {} HostWeightCache::~HostWeightCache() { @@ -1621,7 +1598,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device const void * upload_data = nullptr; size_t upload_size = 0; std::vector transformed; - // Budget 64 MiB of host scratch across phase transitions. The Q6 packer + // Budget 64 MiB of host scratch across first-use materializations. The Q6 packer // writes every output byte, and the cache lock covers packing and upload. std::vector & output = use_scratch ? scratch->second : transformed; result.status = materialize_weight(materialization_source, upload_data, upload_size, output); @@ -1640,16 +1617,8 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device auto entry = std::make_shared(); entry->length = materialized_length; entry->layout = source.layout; - entry->source_identity = source.identity; - entry->source_generation = source.generation; - entry->source_capacity = source.capacity; - entry->source_offset = source.offset; - entry->source_length = source.length; - entry->source_type = source.source_type; - entry->input_size = source.input_size; - entry->output_size = source.output_size; - entry->source_ne = source.source_ne; - entry->source_nb = source.source_nb; + entry->region_length = source.region_length; + entry->region_type = source.region_type; result.status = allocate_device_buffer(device, materialized_length, entry->buffer); if (!result.status.success()) { return result; @@ -1667,6 +1636,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device } entries_.emplace(key, entry); + report_new_representation(key, *entry); ++stats_.misses; stats_.materialized_bytes += materialized_length; result.lease = HostWeightLease(std::move(entry)); @@ -1683,71 +1653,37 @@ HostWeightCacheStats HostWeightCache::stats() const { return result; } -void HostWeightCache::report_coexisting_representations(const std::vector & leases, - const char * reason) { - struct Allocation { - hrx_buffer_t buffer = nullptr; - std::string layout; - size_t bytes = 0; - uint64_t generation = 0; - }; - struct Region { - uint64_t identity = 0; - size_t offset = 0; - size_t length = 0; - std::vector allocations; - }; - std::map regions; - for (const HostWeightLease & lease : leases) { - if (!lease.valid()) { - continue; - } - const HostWeightLease::Entry & entry = *lease.entry_; - std::ostringstream key; - key << entry.source_identity << ':' << entry.source_capacity << ':' - << entry.source_offset << ':' << entry.source_length << ':' << static_cast(entry.source_type); - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - key << ':' << entry.source_ne[dimension] << ':' << entry.source_nb[dimension]; - } - Region & region = regions[key.str()]; - region.identity = entry.source_identity; - region.offset = entry.source_offset; - region.length = entry.source_length; - const bool already_recorded = std::any_of(region.allocations.begin(), region.allocations.end(), - [&](const Allocation & allocation) { return allocation.buffer == entry.buffer; }); - if (!already_recorded) { - region.allocations.push_back({ entry.buffer, entry.layout, entry.length, entry.source_generation }); - } - } - std::lock_guard lock(mutex_); - for (const auto & [key, region] : regions) { - if (region.allocations.size() < 2) { - continue; - } - std::vector sorted = region.allocations; - std::sort(sorted.begin(), sorted.end(), [](const Allocation & lhs, const Allocation & rhs) { - return std::less{}(lhs.buffer, rhs.buffer); - }); - std::ostringstream signature; - signature << key; - std::ostringstream details; - size_t total_bytes = 0; - for (const Allocation & allocation : sorted) { - signature << ':' << allocation.buffer << ':' << allocation.layout << ':' << allocation.bytes - << ':' << allocation.generation; - details << (total_bytes == 0 ? "" : ", ") << allocation.layout << "@generation=" - << allocation.generation << "=" << allocation.bytes << "B"; - total_bytes += allocation.bytes; - } - if (!reported_duplicate_signatures_.insert(signature.str()).second) { +void HostWeightCache::report_new_representation(const SourceKey & key, const HostWeightLease::Entry & entry) { + bool distinct_layout = false; + size_t total_bytes = 0; + std::ostringstream details; + for (const auto & [resident_key, resident] : entries_) { + const bool same_identity = resident_key.identity == key.identity; + const bool same_generation = resident_key.generation == key.generation; + const bool same_capacity = resident_key.capacity == key.capacity; + const bool same_offset = resident_key.offset == key.offset; + const bool same_length = resident->region_length == entry.region_length; + const bool same_type = resident->region_type == entry.region_type; + const bool same_shape = resident_key.source_ne == key.source_ne; + const bool same_strides = resident_key.source_nb == key.source_nb; + const bool same_region = same_identity && same_generation && same_capacity && same_offset && + same_length && same_type && same_shape && same_strides; + if (!same_region) { continue; } - ++stats_.duplicate_warnings; - GGML_LOG_WARN("HRX weight residency: source=%llu region=[%zu,%zu) layouts={%s} resident_bytes=%zu reason=%s\n", - static_cast(region.identity), - region.offset, region.offset + region.length, details.str().c_str(), total_bytes, - reason == nullptr ? "unspecified" : reason); - } + distinct_layout |= resident_key.layout != key.layout; + details << (total_bytes == 0 ? "" : ", ") << resident_key.layout << "=" << resident->length << "B"; + total_bytes += resident->length; + } + if (!distinct_layout) { + return; + } + ++stats_.duplicate_warnings; + GGML_LOG_WARN("HRX weight residency: source=%llu generation=%llu region=[%zu,%zu) layouts={%s}" + " new_layout=%s new_bytes=%zu resident_bytes=%zu\n", + static_cast(key.identity), static_cast(key.generation), + key.offset, key.offset + entry.region_length, details.str().c_str(), key.layout.c_str(), + entry.length, total_bytes); } void HostWeightCache::fail_cold_acquire_after_for_test(size_t successful_acquires) { @@ -1755,28 +1691,6 @@ void HostWeightCache::fail_cold_acquire_after_for_test(size_t successful_acquire fail_cold_after_ = successful_acquires; } -void HostWeightCache::evict_except(const std::vector & current_demand, - const std::vector & outstanding_leases) { - std::lock_guard lock(mutex_); - for (auto entry = entries_.begin(); entry != entries_.end();) { - const HostWeightLease resident(entry->second); - const bool current = std::any_of(current_demand.begin(), current_demand.end(), - [&](const HostWeightSource & source) { - return resident.matches(source); - }); - const bool outstanding = std::any_of(outstanding_leases.begin(), outstanding_leases.end(), - [&](const HostWeightLease & lease) { - return lease.entry_ == entry->second; - }); - const bool retained = current || outstanding; - if (retained) { - ++entry; - } else { - entry = entries_.erase(entry); - } - } -} - void HostWeightCache::clear() { std::lock_guard lock(mutex_); entries_.clear(); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index cad5bb4da6f8..946f53f5057e 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -12,7 +12,6 @@ #include #include #include -#include #include typedef struct hrx_buffer_s * hrx_buffer_t; @@ -110,7 +109,6 @@ class HostWeightLease { hrx_buffer_t buffer() const; size_t length() const; const std::string & layout() const; - bool matches(const HostWeightSource & source) const; private: struct Entry; @@ -127,6 +125,7 @@ struct HostWeightAcquireResult { bool valid() const { return lease.valid() && status.success(); } }; +// Owns each immutable source representation until explicit clearing or backend teardown. class HostWeightCache { public: HostWeightCache(); @@ -140,10 +139,6 @@ class HostWeightCache { HostTransferManager & transfers, const HostWeightSource & source); HostWeightCacheStats stats() const; - // Report only distinct live device allocations for the same retained source region. - void report_coexisting_representations(const std::vector & leases, const char * reason); - void evict_except(const std::vector & current_demand, - const std::vector & outstanding_leases); void fail_cold_acquire_after_for_test(size_t successful_acquires); // Erases lookup ownership; leases and cumulative diagnostics outlive map entries. void clear(); @@ -176,6 +171,8 @@ class HostWeightCache { size_t operator()(const SourceKey & key) const; }; + void report_new_representation(const SourceKey & key, const HostWeightLease::Entry & entry); + mutable std::mutex mutex_; std::unordered_map, SourceKeyHash> entries_; // Size-keyed scratch is overwritten on each use and is not a source-indexed representation cache. @@ -183,7 +180,6 @@ class HostWeightCache { size_t q6_pack_scratch_bytes_ = 0; HostWeightCacheStats stats_; std::shared_ptr accounting_; - std::unordered_set reported_duplicate_signatures_; size_t fail_cold_after_ = SIZE_MAX; }; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index c4edd68e08ca..817daa15a477 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -106,6 +106,7 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ return result; } result.execution_resources = resources; + // HRX command buffers retain queued buffers and executables even after this uncached owner is released. result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); if (!result.success) { result.status.log("execute uncached HRX command program failed"); @@ -217,32 +218,6 @@ void PreparedCommandProgramCache::clear() { programs_.clear(); } -void PreparedCommandProgramCache::invalidate_obsolete_weights(const std::vector & current_demand) { - std::vector> entries; - { - std::lock_guard lock(mutex_); - entries.reserve(programs_.size()); - for (const auto & program : programs_) { - entries.push_back(program.second); - } - } - for (const auto & entry : entries) { - std::lock_guard entry_lock(entry->mutex); - if (entry->resources == nullptr) { - continue; - } - bool obsolete = false; - for (const HostWeightLease & lease : *entry->resources->prepared.resident_host_weights) { - const bool demanded = std::any_of(current_demand.begin(), current_demand.end(), - [&](const HostWeightSource & source) { return lease.matches(source); }); - obsolete |= !demanded; - } - if (obsolete) { - entry->resources.reset(); - } - } -} - void PreparedCommandProgramCache::record_build() { std::lock_guard lock(mutex_); ++stats_.builds; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index 911cb8c9de47..5a0a5c2fef37 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -58,7 +58,6 @@ class PreparedCommandProgramCache { PreparedCommandProgramCacheStats stats() const; void clear(); - void invalidate_obsolete_weights(const std::vector & current_demand); private: struct Key { diff --git a/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h b/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h deleted file mode 100644 index 1bbcd10ea39b..000000000000 --- a/ggml/src/ggml-hrx/runtime/weight-submission-tracker.h +++ /dev/null @@ -1,82 +0,0 @@ -#pragma once - -#include "runtime/command-program-executor.h" -#include "runtime/host-memory.h" -#include "status.h" - -#include -#include -#include - -typedef struct hrx_event_s * hrx_event_t; -typedef struct hrx_semaphore_s * hrx_semaphore_t; - -namespace ggml::hrx { - -// Keeps submitted bindings alive until their stream timeline point or failure-path event completes. -// The current demand is retained separately so completed work can stay warm for the next call. -class WeightSubmissionTracker { - public: - ~WeightSubmissionTracker(); - - WeightSubmissionTracker(const WeightSubmissionTracker &) = delete; - WeightSubmissionTracker & operator=(const WeightSubmissionTracker &) = delete; - WeightSubmissionTracker() = default; - - Status poll_completed(bool * completed = nullptr); - Status record(hrx_device_t device, hrx_stream_t stream, - std::shared_ptr resources, - const std::vector & demand, bool admitted, bool demand_changed); - void mark_stream_synchronized(); - bool has_obsolete_pending() const; - void fail_completion_record_for_test() { fail_next_record_ = true; } - - const std::vector & current_demand() const { return current_demand_; } - uint64_t revision() const { return revision_; } - bool demand_matches(const std::vector & demand) const; - std::vector merged_demand_for(const std::vector & demand) const; - void release_current_leases() { current_leases_.reset(); } - const std::vector & current_leases() const { - static const std::vector empty; - return current_leases_ != nullptr ? *current_leases_ : empty; - } - std::vector outstanding_leases() const; - - private: - struct RegionKey { - uint64_t identity = 0; - size_t offset = 0; - size_t length = 0; - ggml_type type = GGML_TYPE_COUNT; - std::array ne = {}; - std::array nb = {}; - - bool operator==(const RegionKey & other) const; - }; - struct RegionKeyHash { - size_t operator()(const RegionKey & key) const; - }; - using RegionDemands = std::unordered_map, RegionKeyHash>; - static RegionKey region_of(const HostWeightSource & source); - static RegionDemands group_demands(const std::vector & demand); - void publish_demand(const std::vector & demand); - - struct Pending { - hrx_event_t event = nullptr; - hrx_semaphore_t semaphore = nullptr; - uint64_t value = 0; - std::shared_ptr resources; - bool may_hold_obsolete_demand = false; - }; - - static void release_completion(Pending & pending); - - std::vector current_demand_; - RegionDemands retained_regions_; - std::shared_ptr> current_leases_; - uint64_t revision_ = 0; - std::vector pending_; - bool fail_next_record_ = false; -}; - -} // namespace ggml::hrx diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index f8b3afd9fde9..6051882c80ee 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -660,6 +660,8 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { source.capacity = weight_bytes; source.length = weight_bytes; source.materialized_length = weight_bytes; + source.region_length = weight_bytes; + source.region_type = GGML_TYPE_Q6_K; source.layout = ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; source.source_type = GGML_TYPE_Q6_K; source.input_size = 256; @@ -675,8 +677,7 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { } first.lease = {}; - weights.evict_except({}, {}); - REQUIRE(weights.stats().allocation_count == 0); + REQUIRE(weights.stats().allocation_count == 1); std::fill(host.begin(), host.end(), 0xa5); source.generation = 2; auto second = weights.acquire(context->device->device, context->stream, transfers, source); @@ -686,7 +687,26 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { for (uint8_t byte : downloaded) { REQUIRE(byte == 0xa5); } - REQUIRE(weights.stats().peak_resident_bytes == weight_bytes); + REQUIRE(weights.stats().allocation_count == 2); + REQUIRE(weights.stats().peak_resident_bytes == 2 * weight_bytes); + + source.layout = ggml::hrx::kNativeWeightLayout; + auto native = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(native.valid()); + REQUIRE(weights.stats().duplicate_warnings == 1); + auto alias = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(alias.valid()); + REQUIRE(alias.lease.buffer() == native.lease.buffer()); + REQUIRE(weights.stats().duplicate_warnings == 1); + source.generation = 3; + auto changed_native = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(changed_native.valid()); + REQUIRE(weights.stats().duplicate_warnings == 1); + source.layout = ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + auto changed_packed = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(changed_packed.valid()); + REQUIRE(weights.stats().duplicate_warnings == 2); + REQUIRE(weights.stats().allocation_count == 5); } template diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index e13cb6063d2c..97eb4ec1faf6 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -2939,7 +2939,6 @@ static void run_mutable_host_input_after_writeback_case(bool hrx_host) { // The recorded launch advances once; the queued host-download tail must advance again. REQUIRE(after_writer.value >= before_writer.value + 2); } - REQUIRE(backend_state->weight_submissions.poll_completed().success()); REQUIRE(backend_state->graph_replay_state.has_pending_host_writebacks()); // Finish reading registered host memory without draining the backend's deferred writebacks. REQUIRE(!ggml::hrx::take_status(hrx_stream_synchronize(backend_state->stream))); @@ -2968,99 +2967,7 @@ static void run_mutable_host_input_after_writeback_case(bool hrx_host) { enum class PackedWeightSourceKind { Imported, CpuHost, HrxHost, CoherentSubregion }; -static void run_weight_demand_comparison_case() { - ggml::hrx::WeightSubmissionTracker tracker; - ggml::hrx::HostWeightSource native; - native.identity = 17; - native.generation = 3; - native.capacity = 65536; - native.offset = 4096; - native.length = 32768; - native.region_length = 32768; - native.region_type = GGML_TYPE_F16; - native.materialized_length = 32768; - native.source_type = GGML_TYPE_F16; - native.input_size = 256; - native.output_size = 64; - native.source_ne = { 256, 64, 1, 1 }; - native.source_nb = { 2, 512, 32768, 32768 }; - ggml::hrx::HostWeightSource packed = native; - packed.layout = "packed-test-layout"; - packed.materialized_length = 49152; - ggml::hrx::HostWeightSource other = native; - other.offset = 36864; - - REQUIRE(tracker.demand_matches({})); - REQUIRE(!tracker.demand_matches({ native })); - REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native, other }, true, true).success()); - REQUIRE(tracker.demand_matches({})); - REQUIRE(tracker.demand_matches({ native })); - int aliases[2] = {}; - ggml::hrx::HostWeightSource alias = native; - alias.host_data = &aliases[1]; - REQUIRE(tracker.demand_matches({ native, alias, native })); - REQUIRE(tracker.demand_matches({ other, other })); - REQUIRE(!tracker.demand_matches({ native, packed })); - REQUIRE(!tracker.demand_matches({ packed, native })); - REQUIRE(!tracker.demand_matches({ packed })); - - auto expect_mismatch = [&](const ggml::hrx::HostWeightSource & changed) { - REQUIRE(!tracker.demand_matches({ changed })); - }; - ggml::hrx::HostWeightSource changed = native; - changed.generation++; - expect_mismatch(changed); - changed = native; - changed.capacity++; - expect_mismatch(changed); - changed = native; - changed.length++; - expect_mismatch(changed); - changed = native; - changed.materialized_length++; - expect_mismatch(changed); - changed = native; - changed.source_type = GGML_TYPE_F32; - expect_mismatch(changed); - changed = native; - changed.input_size++; - expect_mismatch(changed); - changed = native; - changed.output_size++; - expect_mismatch(changed); - changed = native; - changed.identity++; - expect_mismatch(changed); - changed = native; - changed.offset++; - expect_mismatch(changed); - changed = native; - changed.region_length++; - expect_mismatch(changed); - changed = native; - changed.region_type = GGML_TYPE_F32; - expect_mismatch(changed); - changed = native; - changed.source_ne[0]++; - expect_mismatch(changed); - changed = native; - changed.source_nb[0]++; - expect_mismatch(changed); - - REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native, packed, other }, true, true).success()); - REQUIRE(tracker.demand_matches({ packed, native, packed, native })); - REQUIRE(tracker.demand_matches({ other, native, packed })); - REQUIRE(!tracker.demand_matches({ native })); - ggml::hrx::HostWeightSource third = packed; - third.layout = "another-packed-layout"; - REQUIRE(!tracker.demand_matches({ native, packed, third })); - ggml::hrx::HostWeightSource changed_other = other; - changed_other.generation++; - REQUIRE(!tracker.demand_matches({ changed_other, native, packed })); - REQUIRE(!tracker.demand_matches({ native, packed, changed_other })); -} - -static void run_weight_reclamation_sequence_case() { +static void run_weight_persistence_sequence_case() { constexpr int64_t input_size = 256; constexpr int64_t output_size = 64; ggml_backend_t cpu = init_cpu_backend(); @@ -3194,7 +3101,7 @@ static void run_weight_reclamation_sequence_case() { [](float value) { return std::fabs(value) > 0.1f; })); auto execute = [&](ggml_cgraph * graph, ggml_tensor * output, const std::vector & expected, - size_t live_allocations = 1) { + size_t live_allocations) { REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, output), expected, 2.0e-2f, 3.0e-2f); @@ -3203,182 +3110,93 @@ static void run_weight_reclamation_sequence_case() { REQUIRE(stats.weight_live_allocations == live_allocations); return stats; }; - const auto native_first = execute(hrx_native, hrx_rows, expected_native); + const auto native_first = execute(hrx_native, hrx_rows, expected_native, 1); REQUIRE(native_first.weight_materializations == 1); - REQUIRE(state->weight_submissions.revision() == 1); - const auto packed_first = execute(hrx_packed, hrx_matmul, expected_packed); + REQUIRE(native_first.weight_duplicate_warnings == 0); + const auto packed_first = execute(hrx_packed, hrx_matmul, expected_packed, 2); REQUIRE(packed_first.weight_materializations == 2); - REQUIRE(packed_first.weight_peak_bytes == - std::max(native_first.weight_live_bytes, packed_first.weight_live_bytes)); - const auto before_native_rebuild = state->graph_programs.stats(); - const auto native_second = execute(hrx_native, hrx_rows, expected_native); - REQUIRE(native_second.weight_materializations == 3); - const auto rebuilt_native_demand = state->graph_programs.stats(); - REQUIRE(rebuilt_native_demand.weight_demand_builds > before_native_rebuild.weight_demand_builds); - const uint64_t warm_revision = state->weight_submissions.revision(); - const auto native_unchanged = execute(hrx_native, hrx_rows, expected_native); - REQUIRE(state->weight_submissions.revision() == warm_revision); - REQUIRE(native_unchanged.weight_materializations == native_second.weight_materializations); - REQUIRE(native_unchanged.weight_uploaded_bytes == native_second.weight_uploaded_bytes); - REQUIRE(native_unchanged.prepared_program_hits > native_second.prepared_program_hits); - REQUIRE(state->graph_programs.stats().weight_demand_hits > rebuilt_native_demand.weight_demand_hits); - hrx_timeline_point_t before_replay = {}; - REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &before_replay))); + REQUIRE(packed_first.weight_live_bytes == 2 * weight.size()); + REQUIRE(packed_first.weight_duplicate_warnings == 1); + for (int cycle = 0; cycle < 3; ++cycle) { + const auto native = execute(hrx_native, hrx_rows, expected_native, 2); + const auto packed = execute(hrx_packed, hrx_matmul, expected_packed, 2); + for (const auto & warm : { native, packed }) { + REQUIRE(warm.weight_materializations == packed_first.weight_materializations); + REQUIRE(warm.weight_uploaded_bytes == packed_first.weight_uploaded_bytes); + REQUIRE(warm.weight_peak_bytes == packed_first.weight_peak_bytes); + REQUIRE(warm.weight_duplicate_warnings == packed_first.weight_duplicate_warnings); + REQUIRE(warm.prepared_program_builds == packed_first.prepared_program_builds); + REQUIRE(warm.prepared_program_hits > packed_first.prepared_program_hits); + } + } + const bool replay_disabled = ggml::hrx::graph_replay_disabled_by_environment(); const auto recorded_native = ggml::hrx::GraphExecutor(*state).execute(*hrx_native); REQUIRE(recorded_native.success()); - hrx_timeline_point_t after_replay = {}; - REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &after_replay))); - const char * replay_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); - const bool replay_disabled = replay_flag != nullptr && std::string(replay_flag) == "1"; REQUIRE(recorded_native.graph_replay_event == (replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : ggml::hrx::HrxGraphReplayEvent::Hit)); - if (!replay_disabled) { - REQUIRE(after_replay.semaphore == before_replay.semaphore); - REQUIRE(after_replay.value == before_replay.value + 1); - } - ggml::hrx::WeightSubmissionTracker isolated_tracker; - auto isolated_resources = std::make_shared(); - REQUIRE(isolated_tracker.record(state->device->device, state->stream, isolated_resources, {}, true, false).success()); - hrx_timeline_point_t after_tracker = {}; - REQUIRE(!ggml::hrx::take_status(hrx_stream_get_timeline_position(state->stream, &after_tracker))); - REQUIRE(after_tracker.semaphore == after_replay.semaphore); - REQUIRE(after_tracker.value == after_replay.value); ggml_backend_synchronize(backend); const auto structural_before = state->prepared_programs.stats(); - const auto structural_native = execute(hrx_native_structural, hrx_rows_structural, expected_native); - REQUIRE(structural_native.weight_materializations == native_second.weight_materializations); - const auto structural_packed = execute(hrx_packed_structural, hrx_matmul_structural, expected_packed); - REQUIRE(structural_packed.weight_materializations == native_second.weight_materializations + 1); - const auto structural_return = execute(hrx_native_structural, hrx_rows_structural, expected_native); - REQUIRE(structural_return.weight_materializations == native_second.weight_materializations + 2); + execute(hrx_native_structural, hrx_rows_structural, expected_native, 2); + execute(hrx_packed_structural, hrx_matmul_structural, expected_packed, 2); + const auto structural_return = execute(hrx_native_structural, hrx_rows_structural, expected_native, 2); + REQUIRE(structural_return.weight_materializations == packed_first.weight_materializations); + REQUIRE(structural_return.weight_uploaded_bytes == packed_first.weight_uploaded_bytes); REQUIRE(state->prepared_programs.stats().builds > structural_before.builds); - const auto distinct_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); + const auto distinct_second = execute(hrx_second_native, hrx_second_rows, expected_second, 3); REQUIRE(distinct_second.weight_materializations == structural_return.weight_materializations + 1); + REQUIRE(distinct_second.weight_duplicate_warnings == packed_first.weight_duplicate_warnings); REQUIRE(ggml_backend_graph_compute(backend, hrx_empty) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); REQUIRE(hrx_no_weight->n_nodes > 0); - const auto no_weight = execute(hrx_no_weight, hrx_no_weight_output, expected_no_weight, 2); + const auto no_weight = execute(hrx_no_weight, hrx_no_weight_output, expected_no_weight, 3); REQUIRE(no_weight.weight_materializations == distinct_second.weight_materializations); REQUIRE(no_weight.weight_uploaded_bytes == distinct_second.weight_uploaded_bytes); REQUIRE(no_weight.weight_live_bytes == distinct_second.weight_live_bytes); - const auto reacquired_native = execute(hrx_native, hrx_rows, expected_native, 2); - REQUIRE(reacquired_native.weight_materializations == no_weight.weight_materializations); - REQUIRE(reacquired_native.weight_uploaded_bytes == no_weight.weight_uploaded_bytes); - const auto repeated_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); - REQUIRE(repeated_second.weight_materializations == reacquired_native.weight_materializations); - REQUIRE(repeated_second.weight_uploaded_bytes == reacquired_native.weight_uploaded_bytes); - const auto rerecorded_native = ggml::hrx::GraphExecutor(*state).execute(*hrx_native); - REQUIRE(rerecorded_native.success()); - REQUIRE(rerecorded_native.graph_replay_event == (replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : - ggml::hrx::HrxGraphReplayEvent::Hit)); - ggml_backend_synchronize(backend); - const auto packed_with_second = execute(hrx_packed, hrx_matmul, expected_packed, 2); - const auto second_between_packed = execute(hrx_second_native, hrx_second_rows, expected_second, 2); - const auto packed_reused = execute(hrx_packed, hrx_matmul, expected_packed, 2); - REQUIRE(second_between_packed.weight_materializations == packed_with_second.weight_materializations); - REQUIRE(packed_reused.weight_materializations == packed_with_second.weight_materializations); - REQUIRE(packed_reused.weight_uploaded_bytes == packed_with_second.weight_uploaded_bytes); + execute(hrx_native, hrx_rows, expected_native, 3); + execute(hrx_second_native, hrx_second_rows, expected_second, 3); + const auto packed_reused = execute(hrx_packed, hrx_matmul, expected_packed, 3); + REQUIRE(packed_reused.weight_materializations == no_weight.weight_materializations); + REQUIRE(packed_reused.weight_uploaded_bytes == no_weight.weight_uploaded_bytes); + auto verify_two_weights = [&](ggml_cgraph * graph, ggml_tensor * first, ggml_tensor * second) { REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, first), expected_native, 2.0e-2f, 3.0e-2f); require_close(get_f32_tensor(backend, second), expected_third, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats stats = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); - REQUIRE(stats.weight_live_allocations == 3); + REQUIRE(state->host_weights.stats().allocation_count == 3); }; - auto fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second, - size_t warm_allocations) { - const auto before_cold_failure = execute(hrx_packed, hrx_matmul, expected_packed, warm_allocations); - state->host_weights.fail_cold_acquire_after_for_test(1); - const auto before_failure = state->graph_programs.stats(); - const auto retained_before_failure = state->weight_submissions.current_demand(); - REQUIRE(ggml_backend_graph_compute(backend, two_weights) == GGML_STATUS_FAILED); - REQUIRE(state->graph_programs.stats().weight_demand_builds == before_failure.weight_demand_builds); - REQUIRE(state->weight_submissions.demand_matches(retained_before_failure)); + auto fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second) { ggml_backend_synchronize(backend); - ggml_backend_hrx_cache_stats failed = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &failed)); - REQUIRE(failed.weight_live_allocations == 1); - REQUIRE(failed.weight_materializations == before_cold_failure.weight_materializations + 1); - const auto preserved_second = execute(hrx_second_native, hrx_second_rows, expected_second, 1); - REQUIRE(preserved_second.weight_materializations == failed.weight_materializations); - execute(hrx_packed, hrx_matmul, expected_packed, 2); - verify_two_weights(two_weights, first, second); - }; - fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows, 2); - // A layout transition makes the next native acquisition cold without changing the immutable source. - execute(hrx_third_packed, hrx_third_matmul, expected_third_packed, 3); - const auto structural_failure_before = state->prepared_programs.stats(); - fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second, 3); - REQUIRE(state->prepared_programs.stats().builds > structural_failure_before.builds); - verify_two_weights(hrx_two_weights, hrx_rows, hrx_third_rows); - - auto completion_fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second) { - execute(hrx_packed, hrx_matmul, expected_packed, 3); - state->weight_submissions.fail_completion_record_for_test(); - const auto before_failure = state->graph_programs.stats(); - const auto retained_before_failure = state->weight_submissions.current_demand(); + state->graph_programs.clear(); + state->prepared_programs.clear(); + state->host_weights.clear(); + const auto before_failure = execute(hrx_second_native, hrx_second_rows, expected_second, 1); + const auto owner = state->graph_programs.get_or_build( + *hrx_two_weights, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); + REQUIRE(owner.valid()); + state->host_weights.fail_cold_acquire_after_for_test(1); REQUIRE(ggml_backend_graph_compute(backend, two_weights) == GGML_STATUS_FAILED); - REQUIRE(state->graph_programs.stats().weight_demand_builds == before_failure.weight_demand_builds); - REQUIRE(state->weight_submissions.demand_matches(retained_before_failure)); - // The fallback synchronization completed the submitted writes despite the error return. ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, first), expected_native, 2.0e-2f, 3.0e-2f); - require_close(get_f32_tensor(backend, second), expected_third, 2.0e-2f, 3.0e-2f); ggml_backend_hrx_cache_stats failed = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &failed)); REQUIRE(failed.weight_live_allocations == 2); + REQUIRE(failed.weight_materializations == before_failure.weight_materializations + 1); const auto preserved_second = execute(hrx_second_native, hrx_second_rows, expected_second, 2); REQUIRE(preserved_second.weight_materializations == failed.weight_materializations); - execute(hrx_packed, hrx_matmul, expected_packed, 3); verify_two_weights(two_weights, first, second); + REQUIRE(state->host_weights.stats().misses == failed.weight_materializations + 1); }; - completion_fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows); - completion_fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second); - - // Logical demand can survive a discarded prepared owner and its physical allocation. - execute(hrx_packed, hrx_matmul, expected_packed, 3); - const auto packed_lookup = state->graph_programs.get_or_build( - *hrx_packed, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - REQUIRE(packed_lookup.valid()); - const uint64_t retained_revision = state->weight_submissions.revision(); - REQUIRE(packed_lookup.program->admitted_at(retained_revision)); - const auto before_owner_reset = state->host_weights.stats(); - const auto retained_demand = state->weight_submissions.current_demand(); + fail_then_retry(hrx_two_weights, hrx_rows, hrx_third_rows); + fail_then_retry(hrx_two_weights_structural, hrx_two_first, hrx_two_second); + execute(hrx_third_packed, hrx_third_matmul, expected_third_packed, 4); + const auto all_layouts = execute(hrx_packed, hrx_matmul, expected_packed, 5); ggml_backend_synchronize(backend); - state->graph_programs.invalidate_obsolete_weights({}); - state->prepared_programs.invalidate_obsolete_weights({}); - state->weight_submissions.release_current_leases(); - state->host_weights.evict_except({}, {}); - REQUIRE(!packed_lookup.program->has_prepared_program()); - REQUIRE(!packed_lookup.program->admitted_at(retained_revision)); - REQUIRE(state->host_weights.stats().allocation_count == 0); - REQUIRE(state->weight_submissions.demand_matches(retained_demand)); - REQUIRE(state->weight_submissions.revision() == retained_revision); - execute(hrx_packed, hrx_matmul, expected_packed, 1); - REQUIRE(state->host_weights.stats().misses == before_owner_reset.misses + 1); - REQUIRE(state->weight_submissions.revision() == retained_revision); - REQUIRE(packed_lookup.program->admitted_at(retained_revision)); - - // A failed capture of an otherwise warm execution cannot leave its admission reusable. - const auto before_capture_failure = state->graph_programs.stats(); - const auto current_leases = state->weight_submissions.current_leases(); - state->weight_submissions.fail_completion_record_for_test(); - REQUIRE(ggml_backend_graph_compute(backend, hrx_packed) == GGML_STATUS_FAILED); - REQUIRE(!packed_lookup.program->admitted_at(retained_revision)); - REQUIRE(state->weight_submissions.revision() == retained_revision); - REQUIRE(state->graph_programs.stats().weight_demand_builds == before_capture_failure.weight_demand_builds); - REQUIRE(state->graph_programs.stats().weight_demand_hits == before_capture_failure.weight_demand_hits); - REQUIRE(state->weight_submissions.current_leases().size() == current_leases.size()); - for (size_t i = 0; i < current_leases.size(); ++i) { - REQUIRE(state->weight_submissions.current_leases()[i].buffer() == current_leases[i].buffer()); - } - execute(hrx_packed, hrx_matmul, expected_packed, 1); - REQUIRE(state->weight_submissions.revision() == retained_revision); - REQUIRE(packed_lookup.program->admitted_at(retained_revision)); - REQUIRE(state->graph_programs.stats().weight_demand_builds == before_capture_failure.weight_demand_builds + 1); + state->graph_programs.clear(); + state->prepared_programs.clear(); + REQUIRE(state->host_weights.stats().allocation_count == 5); + const auto rebuilt = execute(hrx_packed, hrx_matmul, expected_packed, 5); + REQUIRE(rebuilt.weight_materializations == all_layouts.weight_materializations); + REQUIRE(rebuilt.weight_uploaded_bytes == all_layouts.weight_uploaded_bytes); ggml_backend_buffer_free(hrx_tensors); ggml_backend_buffer_free(cpu_tensors); @@ -3390,7 +3208,7 @@ static void run_weight_reclamation_sequence_case() { ggml_backend_free(cpu); } -static void run_weight_reclamation_overlap_case() { +static void run_weight_persistence_overlap_case() { constexpr int64_t input_size = 256; constexpr int64_t output_size = 64; const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); @@ -3489,8 +3307,7 @@ static void run_weight_reclamation_overlap_case() { REQUIRE(w0_value != nullptr && w1_value != nullptr); REQUIRE(w0_value->id != w1_value->id); - // Preupload the packed layout before closing the gate. Its first coexistence with native - // then occurs across two real submissions, without a cold upload waiting on the gate. + // Materialize both layouts before gating real submissions; cold uploads wait on the stream. REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); auto packed_lookup = state->graph_programs.get_or_build( @@ -3522,12 +3339,11 @@ static void run_weight_reclamation_overlap_case() { REQUIRE(preuploaded.valid()); ggml_backend_hrx_cache_stats cross_before = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cross_before)); - REQUIRE(cross_before.weight_duplicate_warnings == 0); + REQUIRE(cross_before.weight_duplicate_warnings == 1); // The public executors finish their kernel reads before returning. Queue warmed recordings directly // to exercise ownership of genuinely outstanding reads without weakening that completion boundary. struct Submission { std::shared_ptr resources; - std::vector demand; }; auto prepare_submission = [&](ggml_cgraph * graph) { auto lookup = state->graph_programs.get_or_build( @@ -3549,8 +3365,6 @@ static void run_weight_reclamation_overlap_case() { external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); - auto demands = ggml::hrx::collect_command_program_weight_demands(lookup.program->commands(), bindings); - REQUIRE(demands.valid()); const ggml::hrx::CommandProgramExecutionContext context = { state->device->device, state->stream, state->device->architecture.c_str(), &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, @@ -3564,36 +3378,16 @@ static void run_weight_reclamation_overlap_case() { execution.execution_resources->prepared.resident_host_weights); Submission submission; submission.resources = std::move(execution.execution_resources); - for (const auto & demand : demands.weights) { - submission.demand.push_back(demand.source); - } return submission; }; auto queue_submission = [&](Submission & submission) { - const bool changed = !state->weight_submissions.demand_matches(submission.demand); - if (changed) { - state->weight_submissions.release_current_leases(); - } const auto launch = ggml::hrx::take_status(hrx_graph_exec_launch( submission.resources->recorded.exec, state->stream)); if (launch) { return false; } - const auto * weight_list = submission.resources->prepared.resident_host_weights.get(); - const auto capture = state->weight_submissions.record( - state->device->device, state->stream, submission.resources, submission.demand, true, changed); - const bool shares_weights = &state->weight_submissions.current_leases() == weight_list; submission.resources.reset(); - if (!capture.success()) { - return false; - } - const auto & retained = state->weight_submissions.current_demand(); - state->graph_programs.invalidate_obsolete_weights(retained); - state->prepared_programs.invalidate_obsolete_weights(retained); - const auto outstanding = state->weight_submissions.outstanding_leases(); - state->host_weights.evict_except(retained, outstanding); - state->host_weights.report_coexisting_representations(outstanding, "test outstanding submissions"); - return shares_weights; + return !ggml::hrx::take_status(hrx_stream_flush(state->stream)); }; auto cross_native_submission = prepare_submission(hrx_native_graph); auto cross_packed_submission = prepare_submission(hrx_packed_graph); @@ -3619,9 +3413,6 @@ static void run_weight_reclamation_overlap_case() { queue_submission(cross_packed_submission); ggml_backend_hrx_cache_stats cross_queued = {}; const bool cross_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cross_queued); - bool cross_obsolete_completed = true; - const bool cross_poll_ok = state->weight_submissions.poll_completed(&cross_obsolete_completed).success(); - const bool cross_old_demand_pending = state->weight_submissions.has_obsolete_pending(); { std::lock_guard lock(cross_mutex); cross_finished = true; @@ -3629,13 +3420,11 @@ static void run_weight_reclamation_overlap_case() { cross_cv.notify_one(); const bool cross_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cross_gate, 1)); cross_watchdog.join(); - REQUIRE(cross_native && cross_packed && cross_stats_ok && cross_poll_ok && cross_release_ok && !cross_timed_out); - REQUIRE(!cross_obsolete_completed && cross_old_demand_pending); + REQUIRE(cross_native && cross_packed && cross_stats_ok && cross_release_ok && !cross_timed_out); REQUIRE(cross_queued.weight_live_allocations == 2); - REQUIRE(cross_queued.weight_duplicate_warnings == cross_before.weight_duplicate_warnings + 1); + REQUIRE(cross_queued.weight_duplicate_warnings == cross_before.weight_duplicate_warnings); preuploaded.lease = {}; ggml_backend_synchronize(backend); - REQUIRE(!state->weight_submissions.has_obsolete_pending()); require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); hrx_semaphore_release(cross_gate); @@ -3674,7 +3463,7 @@ static void run_weight_reclamation_overlap_case() { queue_submission(cold_native_submission); state->host_weights.fail_cold_acquire_after_for_test(0); const bool cold_failed = - ggml_backend_graph_compute_async(backend, hrx_packed_graph) == GGML_STATUS_FAILED; + ggml_backend_graph_compute_async(backend, hrx_disjoint_graph) == GGML_STATUS_FAILED; ggml_backend_hrx_cache_stats cold_pending = {}; const bool cold_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cold_pending); { @@ -3685,7 +3474,7 @@ static void run_weight_reclamation_overlap_case() { const bool cold_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cold_gate, 1)); cold_watchdog.join(); REQUIRE(native_before_failure && cold_failed && cold_stats_ok && cold_release_ok && !cold_timed_out); - REQUIRE(cold_pending.weight_live_allocations == 1); + REQUIRE(cold_pending.weight_live_allocations == 2); ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); hrx_semaphore_release(cold_gate); @@ -3743,18 +3532,18 @@ static void run_weight_reclamation_overlap_case() { ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats retired = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &retired)); - REQUIRE(retired.weight_live_allocations == 1); - REQUIRE(retired.weight_live_bytes < queued.weight_live_bytes); - REQUIRE(retired.weight_duplicate_warnings == queued.weight_duplicate_warnings); + ggml_backend_hrx_cache_stats completed = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &completed)); + REQUIRE(completed.weight_live_allocations == 2); + REQUIRE(completed.weight_live_bytes == queued.weight_live_bytes); + REQUIRE(completed.weight_duplicate_warnings == queued.weight_duplicate_warnings); REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats reacquired = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &reacquired)); - REQUIRE(reacquired.weight_materializations == retired.weight_materializations + 1); - REQUIRE(reacquired.weight_live_allocations == 1); + ggml_backend_hrx_cache_stats reused = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &reused)); + REQUIRE(reused.weight_materializations == completed.weight_materializations); + REQUIRE(reused.weight_live_allocations == 2); REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); @@ -3773,16 +3562,16 @@ static void run_weight_reclamation_overlap_case() { REQUIRE(mixed_reused.weight_live_allocations == 3); REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); - ggml_backend_hrx_cache_stats native_narrowed = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &native_narrowed)); - REQUIRE(native_narrowed.weight_live_allocations == 2); + ggml_backend_hrx_cache_stats native_reused = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &native_reused)); + REQUIRE(native_reused.weight_live_allocations == 3); REQUIRE(ggml_backend_graph_compute(backend, hrx_disjoint_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); require_close(get_f32_tensor(backend, hrx_disjoint_rows), expected_disjoint, 2.0e-2f, 3.0e-2f); ggml_backend_hrx_cache_stats disjoint_reused = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &disjoint_reused)); - REQUIRE(disjoint_reused.weight_materializations == native_narrowed.weight_materializations); - REQUIRE(disjoint_reused.weight_live_allocations == 2); + REQUIRE(disjoint_reused.weight_materializations == native_reused.weight_materializations); + REQUIRE(disjoint_reused.weight_live_allocations == 3); hrx_semaphore_release(gate); ggml_backend_buffer_free(hrx_tensors); @@ -3796,7 +3585,7 @@ static void run_weight_reclamation_overlap_case() { restore_environment_value("GGML_HRX_DISABLE_GRAPH_REPLAY", had_flag, old_flag_value); } -static void run_failed_submission_weight_lifetime_case() { +static void run_failed_submission_weight_lifetime_case(bool uncached = false) { constexpr int64_t input_size = 256; constexpr int64_t output_size = 64; const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); @@ -3838,6 +3627,11 @@ static void run_failed_submission_weight_lifetime_case() { REQUIRE(imported != nullptr); REQUIRE(ggml_backend_tensor_alloc(imported, hrx_weight, weight.data()) == GGML_STATUS_SUCCESS); ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_buffer_t input_host = ggml_backend_buft_alloc_buffer( + ggml_backend_cpu_buffer_type(), ggml_nbytes(hrx_input)); + REQUIRE(input_host != nullptr); + REQUIRE(ggml_backend_tensor_alloc(input_host, hrx_input, ggml_backend_buffer_get_base(input_host)) == + GGML_STATUS_SUCCESS); ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); @@ -3893,11 +3687,6 @@ static void run_failed_submission_weight_lifetime_case() { REQUIRE(native_source.layout == ggml::hrx::kNativeWeightLayout); REQUIRE(packed_source.layout != native_source.layout); REQUIRE(packed_source.materialized_length > 0); - ggml::hrx::WeightSubmissionTracker tracker; - REQUIRE(tracker.record(nullptr, nullptr, nullptr, { native_source }, true, true).success()); - REQUIRE(tracker.demand_matches({ native_source })); - REQUIRE(!tracker.demand_matches({ packed_source })); - auto preuploaded = state->host_weights.acquire(state->device->device, state->stream, state->host_transfers, packed_source); REQUIRE(preuploaded.valid()); @@ -3908,9 +3697,13 @@ static void run_failed_submission_weight_lifetime_case() { &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, }; - auto execution = packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); + auto execution = uncached ? + state->prepared_programs.execute_with_result(execution_context, 0, packed_lookup.program->command_shape_hash(), + packed_lookup.program->commands(), packed_bindings) : + packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); REQUIRE(execution.success); - REQUIRE(execution.execution_resources->recorded.valid()); + REQUIRE(!execution.execution_resources->prepared.host_staging.empty()); + REQUIRE(execution.execution_resources->recorded.valid() == !uncached); ggml_backend_synchronize(backend); ggml_backend_tensor_set(hrx_output, poisoned.data(), 0, poisoned.size() * sizeof(float)); // The invalid final command fails after the real weight reader is queued, before the completion boundary. @@ -3933,20 +3726,16 @@ static void run_failed_submission_weight_lifetime_case() { const bool execution_failed = !ggml::hrx::execute_prepared_command_program( execution_context, execution.execution_resources->prepared); std::weak_ptr sole_resource = execution.execution_resources; - const bool has_execution_resources = execution.execution_resources != nullptr; - const bool record_ok = has_execution_resources && tracker.record(state->device->device, state->stream, - execution.execution_resources, { packed_source }, false, true).success(); - state->graph_programs.invalidate_obsolete_weights(tracker.current_demand()); + const bool flush_ok = !ggml::hrx::take_status(hrx_stream_flush(state->stream)); execution.execution_resources.reset(); - state->host_weights.evict_except(tracker.current_demand(), {}); + if (uncached) { + state->graph_programs.clear(); + state->kernel_executables.clear(); + } const auto pending_stats = state->host_weights.stats(); - bool completed = true; - const bool poll_ok = tracker.poll_completed(&completed).success(); bool stream_complete = true; const bool query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); - const bool demand_unchanged = tracker.demand_matches({ native_source }); - const bool tracker_only_owner = sole_resource.use_count() == 1; - const bool resource_retained = tracker.outstanding_leases().size() == 1; + const bool cache_only_owner = sole_resource.use_count() == 1; { std::lock_guard lock(gate_mutex); finished = true; @@ -3954,26 +3743,29 @@ static void run_failed_submission_weight_lifetime_case() { gate_cv.notify_one(); const bool signal_ok = !ggml::hrx::take_status(hrx_semaphore_signal(gate, 1)); watchdog.join(); - REQUIRE(execution_failed && record_ok && poll_ok && query_ok && signal_ok && !timed_out); - REQUIRE(!completed && !stream_complete && demand_unchanged && tracker_only_owner && resource_retained); + REQUIRE(execution_failed && flush_ok && query_ok && signal_ok && !timed_out); + REQUIRE(!stream_complete); + REQUIRE(cache_only_owner == !uncached); + REQUIRE(sole_resource.expired() == uncached); REQUIRE(pending_stats.allocation_count == 1); REQUIRE(pending_stats.resident_bytes == packed_source.materialized_length); REQUIRE(!ggml::hrx::take_status(hrx_stream_synchronize(state->stream))); require_close(get_f32_tensor(backend, hrx_output), expected, 2.0e-2f, 3.0e-2f); - bool retired = false; - REQUIRE(tracker.poll_completed(&retired).success()); - REQUIRE(retired); - state->host_weights.evict_except(tracker.current_demand(), tracker.outstanding_leases()); - const auto reclaimed = state->host_weights.stats(); + REQUIRE(sole_resource.expired() == uncached); + const auto completed = state->host_weights.stats(); + REQUIRE(completed.allocation_count == 1); + REQUIRE(completed.resident_bytes == packed_source.materialized_length); + state->graph_programs.clear(); REQUIRE(sole_resource.expired()); - REQUIRE(tracker.demand_matches({ native_source })); - REQUIRE(reclaimed.allocation_count == 0); - REQUIRE(reclaimed.resident_bytes == 0); + REQUIRE(state->host_weights.stats().allocation_count == 1); + state->host_weights.clear(); + REQUIRE(state->host_weights.stats().allocation_count == 0); hrx_semaphore_release(gate); ggml_backend_buffer_free(hrx_tensors); ggml_backend_buffer_free(cpu_tensors); ggml_backend_buffer_free(imported); + ggml_backend_buffer_free(input_host); ggml_free(hrx_ctx); ggml_free(cpu_ctx); ggml_backend_free(backend); @@ -4006,6 +3798,15 @@ static void run_imported_packed_weight_case(ggml_type weight_type, ggml_tensor * hrx_other_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kInputSize, 1); ggml_tensor * cpu_norm_weight = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, kInputSize); ggml_tensor * hrx_norm_weight = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, kInputSize); + ggml_tensor * cpu_prefill_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kInputSize, 32); + ggml_tensor * hrx_prefill_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kInputSize, 32); + ggml_tensor * cpu_prefill_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_prefill_input); + ggml_tensor * hrx_prefill_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_prefill_input); + ggml_cgraph * cpu_prefill = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_prefill = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(cpu_prefill, cpu_prefill_output); + ggml_build_forward_expand(hrx_prefill, hrx_prefill_output); + hrx_prefill->uid = 951004; ggml_tensor * cpu_native_indices = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); ggml_tensor * hrx_native_indices = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); ggml_tensor * cpu_decode_rhs = build_rmsnorm_mul_graph(cpu_ctx, cpu_decode_input, cpu_norm_weight); @@ -4092,6 +3893,9 @@ static void run_imported_packed_weight_case(ggml_type weight_type, ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); REQUIRE(hrx_buffer != nullptr && cpu_buffer != nullptr); ggml_backend_tensor_set(cpu_weight, source.data(), 0, source.size()); + const std::vector prefill_input = make_pattern_f32(kInputSize * 32, 89, 0.01f); + set_tensor_pair_bytes(cpu_backend, cpu_prefill_input, hrx_backend, hrx_prefill_input, prefill_input.data(), + prefill_input.size() * sizeof(float)); const std::vector decode_input = make_pattern_f32(kInputSize, 7, 0.01f); const std::array native_indices = { 1, 17, 63 }; set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, decode_input.data(), @@ -4406,16 +4210,44 @@ static void run_imported_packed_weight_case(ggml_type weight_type, REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); const uint64_t native_materializations = coherent_subregion ? 1 : 2; REQUIRE(stats.weight_materializations == native_materializations); - REQUIRE(stats.weight_live_allocations == (coherent_subregion ? 0 : 1)); + REQUIRE(stats.weight_live_allocations == native_materializations); const uint64_t uploaded_bytes = stats.weight_uploaded_bytes; compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); - REQUIRE(stats.weight_materializations == 2 * native_materializations); - REQUIRE(stats.weight_live_allocations == (coherent_subregion ? 0 : 1)); - REQUIRE(stats.weight_uploaded_bytes == 2 * uploaded_bytes); + REQUIRE(stats.weight_materializations == native_materializations); + REQUIRE(stats.weight_live_allocations == native_materializations); + REQUIRE(stats.weight_uploaded_bytes == uploaded_bytes); + + const auto warmed_layouts = stats; + auto prefill_lookup = backend_state->graph_programs.get_or_build( + *hrx_prefill, ggml::hrx::get_qwen_kernel_corpus(), backend_state->device->architecture); + REQUIRE(prefill_lookup.valid()); + const auto * prefill_weight = prefill_lookup.program->graph().values().find_tensor(hrx_weight); + REQUIRE(prefill_weight != nullptr); + bool native_prefill = false; + for (const auto & command : prefill_lookup.program->commands().commands) { + for (const auto & binding : command.bindings) { + if (binding.value == prefill_weight->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + native_prefill = true; + } + } + } + REQUIRE(native_prefill); + for (int cycle = 0; cycle < 3; ++cycle) { + compute_and_compare(cpu_prefill, hrx_prefill, cpu_prefill_output, hrx_prefill_output); + compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == warmed_layouts.weight_materializations); + REQUIRE(stats.weight_uploaded_bytes == warmed_layouts.weight_uploaded_bytes); + REQUIRE(stats.weight_live_allocations == warmed_layouts.weight_live_allocations); + REQUIRE(stats.weight_live_bytes == warmed_layouts.weight_live_bytes); + REQUIRE(stats.weight_peak_bytes == warmed_layouts.weight_peak_bytes); + REQUIRE(stats.weight_duplicate_warnings == warmed_layouts.weight_duplicate_warnings); + } + REQUIRE(stats.weight_duplicate_warnings == 1); - // The preceding native consumer may retire packed residency; restore packed demand before testing input rebinding. compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); ggml_backend_hrx_cache_stats before_offset_change = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_offset_change)); @@ -4497,7 +4329,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, backend_state->graph_programs.clear(); backend_state->prepared_programs.clear(); lookup = {}; - backend_state->weight_submissions.release_current_leases(); backend_state->host_weights.clear(); REQUIRE(backend_state->host_weights.stats().allocation_count == 1); recorded = {}; @@ -4641,8 +4472,7 @@ static void run_mixed_weight_layout_program_case(ggml_type weight_type) { REQUIRE(repeated.duplicate_warnings == first.duplicate_warnings); const auto repeated_graph_stats = state->graph_programs.stats(); REQUIRE(repeated_graph_stats.prepared_program_hits > 0); - REQUIRE(repeated_graph_stats.weight_demand_builds >= 1); - REQUIRE(repeated_graph_stats.weight_demand_hits >= 1); + REQUIRE(repeated_graph_stats.prepared_program_builds >= 1); ggml_backend_hrx_cache_stats public_repeated = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &public_repeated)); REQUIRE(public_repeated.weight_duplicate_warnings == public_first.weight_duplicate_warnings); @@ -4931,11 +4761,11 @@ static void run_native_cpu_weight_during_import_layout_transition_case(ggml_type compute_and_compare(cpu_native, hrx_native, cpu_imported_rows, hrx_imported_rows); auto native = state->host_weights.stats(); REQUIRE(native.misses == first.misses + 1); - REQUIRE(native.allocation_count == 2); + REQUIRE(native.allocation_count == 3); compute_and_compare(cpu_packed, hrx_packed, cpu_imported_matmul, hrx_imported_matmul); auto packed_again = state->host_weights.stats(); - REQUIRE(packed_again.misses == native.misses + 1); - REQUIRE(packed_again.allocation_count == 2); + REQUIRE(packed_again.misses == native.misses); + REQUIRE(packed_again.allocation_count == 3); ggml_backend_buffer_free(hrx_tensors); ggml_backend_buffer_free(cpu_tensors); @@ -9269,10 +9099,10 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("imported_native_weight", [] { run_imported_native_weight_case(); }); suite.device_case("mutable_host_input_after_writeback.cpu", [] { run_mutable_host_input_after_writeback_case(false); }); suite.device_case("mutable_host_input_after_writeback.hrx", [] { run_mutable_host_input_after_writeback_case(true); }); - suite.host_case("weight_demand_comparison", [] { run_weight_demand_comparison_case(); }); - suite.device_case("weight_reclamation_sequence", [] { run_weight_reclamation_sequence_case(); }); - suite.device_case("weight_reclamation_overlap", [] { run_weight_reclamation_overlap_case(); }); + suite.device_case("weight_persistence_sequence", [] { run_weight_persistence_sequence_case(); }); + suite.device_case("weight_persistence_overlap", [] { run_weight_persistence_overlap_case(); }); suite.device_case("failed_submission_weight_lifetime", [] { run_failed_submission_weight_lifetime_case(); }); + suite.device_case("failed_submission_weight_lifetime.uncached", [] { run_failed_submission_weight_lifetime_case(true); }); suite.device_case("imported_packed_weight.q4_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q4_K); }); suite.device_case("imported_packed_weight.q6_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q6_K); }); suite.device_case("mixed_weight_layout_program.q4_k", [] { run_mixed_weight_layout_program_case(GGML_TYPE_Q4_K); }); From 4809a914bf38c2fb2b4544a790ddac8d0dd95352 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 03:26:41 +0000 Subject: [PATCH 05/18] hrx: key weight residency by produced bytes and trim warm bindings --- ggml/src/ggml-hrx/README.md | 4 +- .../dispatch/command-program-bindings.cpp | 22 +---- .../dispatch/command-program-bindings.h | 7 +- .../runtime/command-program-executor.cpp | 6 +- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 13 +-- ggml/src/ggml-hrx/runtime/host-memory.cpp | 27 ++--- ggml/src/ggml-hrx/runtime/host-memory.h | 10 +- tests/hrx-backend-test.cpp | 14 +-- tests/hrx-buffer-test.cpp | 99 +++++++++++++++---- tests/hrx-ops-test.cpp | 72 +++++--------- 10 files changed, 127 insertions(+), 147 deletions(-) diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index 3e98c3b60dfe..0c7774e8f688 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -8,9 +8,9 @@ Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx Live weight bytes count physical allocations retained by the cache, a prepared program, or a recorded graph. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. -Each immutable source region is materialized only when a command first requests its layout. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. A warm graph reuses its prepared bindings without rebuilding per-weight demand; mutable inputs still rebind and upload normally. +Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. A warm graph reuses its prepared bindings without rebuilding per-weight demand; mutable inputs still rebind and upload normally. -Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. +Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning groups exact source byte ranges; it does not detect partial overlap or infer a whole tensor from a transformed subrange. It names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. Keep borrowed weight bytes unchanged and alive while HRX can use them. A foreign CPU host buffer and its storage must stay alive for the lifetime of the HRX backend that caches it. Its identity depends on its wrapper and base addresses; destroying and recreating a buffer at those addresses while that backend lives can reuse old residency. Training and in-place weight updates are outside this path. Ordinary mutable inputs and outputs retain their host staging, writeback and synchronization behavior. diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 191a8dc8b558..29a3193b2d06 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -14,18 +14,13 @@ static void mix_hash(uint64_t & hash, uint64_t value) { } // namespace CommandProgramBinding make_command_program_binding( - ValueId value, const ValueBufferBinding & buffer, ggml_type type, - const std::array & ne, const std::array & nb, - bool empty_value) { + ValueId value, const ValueBufferBinding & buffer, bool empty_value) { CommandProgramBinding binding = { value, buffer.buffer, buffer.offset, buffer.length, buffer.identity, buffer.generation, buffer.capacity, buffer.host_data, buffer.weight, empty_value }; binding.source_host_data = buffer.source_host_data; binding.source_offset = buffer.source_offset; binding.source_capacity = buffer.source_capacity; - binding.source_type = type; - binding.source_ne = ne; - binding.source_nb = nb; return binding; } @@ -43,9 +38,7 @@ CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & v result.status.log("external value %d is not bound", id.value); continue; } - CommandProgramBinding binding = make_command_program_binding( - value->id, *buffer, value->type, value->ne, value->nb, value->byte_count == 0); - bindings.push_back(binding); + bindings.push_back(make_command_program_binding(value->id, *buffer, value->byte_count == 0)); } return from_bindings(std::move(bindings), result.status); } @@ -108,11 +101,6 @@ CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBin mix_hash(hash, binding.source_host_data != nullptr ? 1 : 0); mix_hash(hash, static_cast(binding.source_offset)); mix_hash(hash, static_cast(binding.source_capacity)); - mix_hash(hash, static_cast(binding.source_type)); - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - mix_hash(hash, static_cast(binding.source_ne[dimension])); - mix_hash(hash, static_cast(binding.source_nb[dimension])); - } } mix_hash(hash, static_cast(bindings.bindings().size())); return { hash }; @@ -128,11 +116,7 @@ CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const Com << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0) << "|source_host=" << (binding.source_host_data != nullptr ? 1 : 0) - << "|source_offset=" << binding.source_offset << "|source_capacity=" << binding.source_capacity - << "|source_type=" << static_cast(binding.source_type); - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - out << "|ne=" << binding.source_ne[dimension] << "|nb=" << binding.source_nb[dimension]; - } + << "|source_offset=" << binding.source_offset << "|source_capacity=" << binding.source_capacity; } return { out.str() }; } diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index 5356eae4771c..1358bb9cf9fc 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -6,7 +6,6 @@ #include #include -#include #include #include @@ -29,9 +28,6 @@ struct CommandProgramBinding { const void * source_host_data = nullptr; size_t source_offset = 0; size_t source_capacity = 0; - ggml_type source_type = GGML_TYPE_COUNT; - std::array source_ne = {}; - std::array source_nb = {}; // Materialized bindings are selected by the requesting command, not just by ValueId. std::string layout = kNativeWeightLayout; ggml_type layout_source_type = GGML_TYPE_COUNT; @@ -45,8 +41,7 @@ struct CommandProgramBinding { }; CommandProgramBinding make_command_program_binding( - ValueId value, const ValueBufferBinding & buffer, ggml_type type, - const std::array & ne, const std::array & nb, bool empty_value); + ValueId value, const ValueBufferBinding & buffer, bool empty_value); struct CommandProgramBindingsFingerprint { std::string value; diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 124f0f9183b0..ca6b7620f0e0 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -1853,15 +1853,11 @@ CommandProgramWeightDemands bind_command_program_weight_demands(const CommandPro demand.source.capacity = binding->source_capacity; demand.source.offset = binding->source_offset; demand.source.length = transformed ? request.source_length : binding->length; - demand.source.region_length = binding->length; - demand.source.region_type = binding->source_type; demand.source.materialized_length = transformed ? request.length : binding->length; demand.source.layout = request.layout; - demand.source.source_type = transformed ? request.source_type : binding->source_type; + demand.source.source_type = transformed ? request.source_type : GGML_TYPE_COUNT; demand.source.input_size = request.input_size; demand.source.output_size = request.output_size; - demand.source.source_ne = binding->source_ne; - demand.source.source_nb = binding->source_nb; result.weights.push_back(std::move(demand)); } return result; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index c9895ae7644f..044010b5e011 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -53,21 +53,12 @@ CommandProgramBindings GraphExecutor::bind_external_value_buffers(const GraphPro bindings.reserve(match.external_bindings.size()); for (const GraphProgramExternalBinding & external : match.external_bindings) { ValueBufferBinding value_binding; - CommandProgramBinding binding; - binding.value = external.value; if (ggml_backend_hrx_resolve_value_buffer(external.tensor, value_binding)) { - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = external.tensor->ne[dimension]; - nb[dimension] = external.tensor->nb[dimension]; - } - binding = make_command_program_binding(external.value, value_binding, external.tensor->type, ne, nb, - ggml_nbytes(external.tensor) == 0); + bindings.push_back(make_command_program_binding(external.value, value_binding, + ggml_nbytes(external.tensor) == 0)); } else { status.log("external value %d is not bound", external.value.value); } - bindings.push_back(binding); } return CommandProgramBindings::from_bindings(std::move(bindings), status); } diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 1900468b45cc..156dfe452203 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -1452,8 +1452,6 @@ struct HostWeightLease::Entry { hrx_buffer_t buffer = nullptr; size_t length = 0; std::string layout; - size_t region_length = 0; - ggml_type region_type = GGML_TYPE_COUNT; std::shared_ptr accounting; }; @@ -1500,10 +1498,6 @@ size_t HostWeightCache::SourceKeyHash::operator()(const SourceKey & key) const { mix(static_cast(key.source_type)); mix(static_cast(key.input_size)); mix(static_cast(key.output_size)); - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - mix(static_cast(key.source_ne[dimension])); - mix(static_cast(key.source_nb[dimension])); - } return static_cast(hash); } @@ -1535,9 +1529,12 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device return result; } - const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, - source.length, materialized_length, source.layout, source.source_type, - source.input_size, source.output_size, source.source_ne, source.source_nb }; + // Native residency copies this byte range verbatim. Only transformed layouts interpret its type and shape. + const bool native_layout = source.layout == kNativeWeightLayout; + const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, + source.length, materialized_length, source.layout, + native_layout ? GGML_TYPE_COUNT : source.source_type, + native_layout ? 0 : source.input_size, native_layout ? 0 : source.output_size }; std::lock_guard lock(mutex_); const auto found = entries_.find(key); if (found != entries_.end()) { @@ -1617,8 +1614,6 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device auto entry = std::make_shared(); entry->length = materialized_length; entry->layout = source.layout; - entry->region_length = source.region_length; - entry->region_type = source.region_type; result.status = allocate_device_buffer(device, materialized_length, entry->buffer); if (!result.status.success()) { return result; @@ -1654,6 +1649,7 @@ HostWeightCacheStats HostWeightCache::stats() const { } void HostWeightCache::report_new_representation(const SourceKey & key, const HostWeightLease::Entry & entry) { + // Report representations of the same source byte range; partially overlapping ranges are independent. bool distinct_layout = false; size_t total_bytes = 0; std::ostringstream details; @@ -1662,12 +1658,9 @@ void HostWeightCache::report_new_representation(const SourceKey & key, const Hos const bool same_generation = resident_key.generation == key.generation; const bool same_capacity = resident_key.capacity == key.capacity; const bool same_offset = resident_key.offset == key.offset; - const bool same_length = resident->region_length == entry.region_length; - const bool same_type = resident->region_type == entry.region_type; - const bool same_shape = resident_key.source_ne == key.source_ne; - const bool same_strides = resident_key.source_nb == key.source_nb; + const bool same_length = resident_key.length == key.length; const bool same_region = same_identity && same_generation && same_capacity && same_offset && - same_length && same_type && same_shape && same_strides; + same_length; if (!same_region) { continue; } @@ -1682,7 +1675,7 @@ void HostWeightCache::report_new_representation(const SourceKey & key, const Hos GGML_LOG_WARN("HRX weight residency: source=%llu generation=%llu region=[%zu,%zu) layouts={%s}" " new_layout=%s new_bytes=%zu resident_bytes=%zu\n", static_cast(key.identity), static_cast(key.generation), - key.offset, key.offset + entry.region_length, details.str().c_str(), key.layout.c_str(), + key.offset, key.offset + key.length, details.str().c_str(), key.layout.c_str(), entry.length, total_bytes); } diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 946f53f5057e..7cd4197228fd 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -6,7 +6,6 @@ #include "status.h" #include -#include #include #include #include @@ -77,16 +76,11 @@ struct HostWeightSource { size_t capacity = 0; size_t offset = 0; size_t length = 0; - // Native extent and type identify the logical weight across its requested layouts. - size_t region_length = 0; - ggml_type region_type = GGML_TYPE_COUNT; size_t materialized_length = 0; std::string layout = kNativeWeightLayout; ggml_type source_type = GGML_TYPE_COUNT; int64_t input_size = 0; int64_t output_size = 0; - std::array source_ne = {}; - std::array source_nb = {}; }; struct HostWeightCacheStats { @@ -155,15 +149,13 @@ class HostWeightCache { ggml_type source_type = GGML_TYPE_COUNT; int64_t input_size = 0; int64_t output_size = 0; - std::array source_ne = {}; - std::array source_nb = {}; bool operator==(const SourceKey & other) const { return identity == other.identity && generation == other.generation && capacity == other.capacity && offset == other.offset && length == other.length && materialized_length == other.materialized_length && layout == other.layout && source_type == other.source_type && input_size == other.input_size && - output_size == other.output_size && source_ne == other.source_ne && source_nb == other.source_nb; + output_size == other.output_size; } }; diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index be3f75cd68ff..65e399e0eb0f 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -11747,7 +11747,7 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { const auto * value = lookup.program->graph().values().find(external.value); REQUIRE(value != nullptr); external_bindings.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, value->type, value->ne, value->nb, value->byte_count == 0)); + external.value, buffer, value->byte_count == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); REQUIRE(bindings.valid()); @@ -12319,9 +12319,6 @@ static void run_command_weight_request_checks() { binding.weight = true; binding.source_host_data = source.data(); binding.source_capacity = source.size(); - binding.source_type = GGML_TYPE_F32; - binding.source_ne = { 8, 2, 1, 1 }; - binding.source_nb = { 4, 32, 64, 64 }; CommandBinding native = { "first-native", ValueId(1) }; CommandBinding packed = native; @@ -12405,9 +12402,6 @@ static void run_command_weight_request_checks() { next.source_host_data = source.data() + 16; next.source_offset = 32; next.source_capacity = source.size() - 16; - next.source_type = GGML_TYPE_F16; - next.source_ne = { 16, 4, 1, 1 }; - next.source_nb = { 2, 32, 128, 128 }; } const auto rebound = bind_command_program_weight_demands( requests, CommandProgramBindings::from_bindings(current)); @@ -12419,15 +12413,11 @@ static void run_command_weight_request_checks() { REQUIRE(demand.source.generation == 40); REQUIRE(demand.source.offset == 32); REQUIRE(demand.source.capacity == source.size() - 16); - REQUIRE(demand.source.region_length == 128); - REQUIRE(demand.source.region_type == GGML_TYPE_F16); - REQUIRE(demand.source.source_ne == current[0].source_ne); - REQUIRE(demand.source.source_nb == current[0].source_nb); REQUIRE(demand.requires_materialization); } REQUIRE(rebound.weights[0].source.length == 128); REQUIRE(rebound.weights[0].source.materialized_length == 128); - REQUIRE(rebound.weights[0].source.source_type == GGML_TYPE_F16); + REQUIRE(rebound.weights[0].source.source_type == GGML_TYPE_COUNT); REQUIRE(rebound.weights[1].source.length == 64); REQUIRE(rebound.weights[1].source.source_type == GGML_TYPE_F32); for (auto & next : current) { diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index 6051882c80ee..c136808c8845 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -616,8 +616,6 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { source.generation = 3; source.offset = 16; source.source_type = GGML_TYPE_F32; - source.source_ne = { 8, 1, 1, 1 }; - source.source_nb = { 4, 32, 32, 32 }; auto region = weights.acquire(context->device->device, context->stream, transfers, source); auto alias = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(region.valid() && alias.valid()); @@ -625,27 +623,28 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { source.source_type = GGML_TYPE_I32; auto other_type = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(other_type.valid()); - REQUIRE(other_type.lease.buffer() != region.lease.buffer()); - source.source_type = GGML_TYPE_F32; - source.source_ne = { 4, 2, 1, 1 }; + REQUIRE(other_type.lease.buffer() == region.lease.buffer()); + // A native copy does not interpret element type or transformation dimensions. + source.input_size = 4; + source.output_size = 2; auto other_shape = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(other_shape.valid()); - REQUIRE(other_shape.lease.buffer() != region.lease.buffer()); - source.source_ne = { 8, 1, 1, 1 }; - source.source_nb = { 4, 64, 64, 64 }; - auto other_stride = weights.acquire(context->device->device, context->stream, transfers, source); - REQUIRE(other_stride.valid()); - REQUIRE(other_stride.lease.buffer() != region.lease.buffer()); - source.source_nb = { 4, 32, 32, 32 }; + REQUIRE(other_shape.lease.buffer() == region.lease.buffer()); + std::array aliased_bytes = {}; + require_hrx_status(hrx_synchronous_d2h(context->device->device, other_shape.lease.buffer(), 0, + aliased_bytes.data(), aliased_bytes.size())); + for (size_t i = 0; i < aliased_bytes.size(); ++i) { + REQUIRE(aliased_bytes[i] == host[source.offset + i]); + } source.offset = 32; auto other_slice = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(other_slice.valid()); REQUIRE(other_slice.lease.buffer() != region.lease.buffer()); weight_stats = weights.stats(); - REQUIRE(weight_stats.hits == 2); - REQUIRE(weight_stats.misses == 8); - REQUIRE(weight_stats.allocation_count == 5); - REQUIRE(weight_stats.peak_resident_bytes == 160); + REQUIRE(weight_stats.hits == 4); + REQUIRE(weight_stats.misses == 5); + REQUIRE(weight_stats.allocation_count == 2); + REQUIRE(weight_stats.peak_resident_bytes == 96); } static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { @@ -660,8 +659,6 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { source.capacity = weight_bytes; source.length = weight_bytes; source.materialized_length = weight_bytes; - source.region_length = weight_bytes; - source.region_type = GGML_TYPE_Q6_K; source.layout = ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; source.source_type = GGML_TYPE_Q6_K; source.input_size = 256; @@ -709,6 +706,67 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { REQUIRE(weights.stats().allocation_count == 5); } +static void run_packed_weight_identity_checks(ggml_backend_hrx_context * context) { + constexpr size_t source_bytes = 128 * 210; + std::vector host(source_bytes); + for (size_t i = 0; i < host.size(); ++i) { + host[i] = static_cast((i / 210 + 5 * (i % 210)) % 251); + } + ggml::hrx::HostTransferManager transfers; + ggml::hrx::HostWeightCache weights; + ggml::hrx::HostWeightSource source; + source.host_data = host.data(); + source.identity = 0x5679; + source.generation = 1; + source.capacity = source_bytes; + source.length = source_bytes; + source.materialized_length = source_bytes; + source.layout = ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + source.source_type = GGML_TYPE_Q6_K; + source.input_size = 256; + source.output_size = 128; + auto first = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(first.valid()); + std::vector first_bytes(source_bytes); + require_hrx_status(hrx_synchronous_d2h(context->device->device, first.lease.buffer(), 0, + first_bytes.data(), first_bytes.size())); + + // Equal source/output byte lengths can still require different packed byte orders. + source.input_size = 512; + source.output_size = 64; + auto reshaped = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(reshaped.valid()); + REQUIRE(reshaped.lease.buffer() != first.lease.buffer()); + std::vector reshaped_bytes(source_bytes); + require_hrx_status(hrx_synchronous_d2h(context->device->device, reshaped.lease.buffer(), 0, + reshaped_bytes.data(), reshaped_bytes.size())); + REQUIRE(reshaped_bytes != first_bytes); + auto alias = weights.acquire(context->device->device, context->stream, transfers, source); + REQUIRE(alias.valid()); + REQUIRE(alias.lease.buffer() == reshaped.lease.buffer()); + + auto reject = [&](const ggml::hrx::HostWeightSource & invalid) { + auto result = weights.acquire(context->device->device, context->stream, transfers, invalid); + REQUIRE(!result.valid()); + }; + auto invalid = source; + invalid.source_type = GGML_TYPE_Q5_K; + reject(invalid); + invalid = source; + ++invalid.input_size; + reject(invalid); + invalid = source; + --invalid.output_size; + reject(invalid); + invalid = source; + --invalid.length; + reject(invalid); + invalid = source; + --invalid.materialized_length; + reject(invalid); + REQUIRE(weights.stats().allocation_count == 2); +} + template static void run_with_hrx_backend(F && f) { ggml_backend_t backend = ggml_backend_hrx_init(0); @@ -754,6 +812,11 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_q6_pack_scratch_checks(context); }); }); + suite.device_case("packed_weight_identity", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_packed_weight_identity_checks(context); + }); + }); } int main(int argc, char ** argv) { diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 97eb4ec1faf6..3c50fbcc36c9 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -2830,16 +2830,26 @@ static void run_imported_native_weight_case() { ggml_tensor * weight_alias = ggml_view_2d(ctx, weight, kInputSize, kOutputSize, weight->nb[1], 0); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kInputSize, 1); ggml_tensor * output = ggml_mul_mat(ctx, weight_alias, input); + ggml_tensor * reshaped_weight = ggml_view_2d( + ctx, weight, kInputSize / 2, kOutputSize * 2, kInputSize / 2 * sizeof(float), 0); + ggml_tensor * reshaped_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kInputSize / 2, 1); + ggml_tensor * reshaped_output = ggml_mul_mat(ctx, reshaped_weight, reshaped_input); REQUIRE(output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); graph->uid = 950001; + ggml_cgraph * reshaped_graph = ggml_new_graph(ctx); + ggml_build_forward_expand(reshaped_graph, reshaped_output); + reshaped_graph->uid = 950002; std::vector weight_data = make_router_weight(kInputSize, kOutputSize); const std::vector input_data = make_router_input(kInputSize, 1); const std::vector expected = router_projection_reference(input_data, weight_data, kInputSize, kOutputSize, 1); + const auto reshaped_input_data = make_router_input(kInputSize / 2, 1); + const auto reshaped_expected = router_projection_reference( + reshaped_input_data, weight_data, kInputSize / 2, kOutputSize * 2, 1); ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( ggml_backend_get_device(backend), weight_data.data(), weight_data.size() * sizeof(float), weight_data.size() * sizeof(float)); @@ -2849,6 +2859,7 @@ static void run_imported_native_weight_case() { ggml_backend_buffer_t transient = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(transient != nullptr); ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(reshaped_input, reshaped_input_data.data(), 0, reshaped_input_data.size() * sizeof(float)); ggml_backend_hrx_imported_region region = {}; REQUIRE(ggml_backend_hrx_get_imported_region(weight, region)); @@ -2869,6 +2880,12 @@ static void run_imported_native_weight_case() { ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); require_close(actual, expected, 1.0e-2f); } + // Graph geometry changes how the same native bytes are interpreted, not their resident representation. + for (int run = 0; run < 2; ++run) { + REQUIRE(ggml_backend_graph_compute(backend, reshaped_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + require_close(get_f32_tensor(backend, reshaped_output), reshaped_expected, 1.0e-2f); + } const auto after = backend_state->host_weights.stats(); REQUIRE(after.misses == before.misses + 1); REQUIRE(after.allocation_count == before.allocation_count + 1); @@ -3319,14 +3336,8 @@ static void run_weight_persistence_overlap_case() { for (const auto & external : packed_match.external_bindings) { ggml::hrx::ValueBufferBinding buffer; REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = external.tensor->ne[dimension]; - nb[dimension] = external.tensor->nb[dimension]; - } packed_external.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + external.value, buffer, ggml_nbytes(external.tensor) == 0)); } auto packed_bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(packed_external)); REQUIRE(packed_bindings.valid()); @@ -3355,14 +3366,8 @@ static void run_weight_persistence_overlap_case() { for (const auto & external : match.external_bindings) { ggml::hrx::ValueBufferBinding buffer; REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = external.tensor->ne[dimension]; - nb[dimension] = external.tensor->nb[dimension]; - } external_bindings.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + external.value, buffer, ggml_nbytes(external.tensor) == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); const ggml::hrx::CommandProgramExecutionContext context = { @@ -3655,14 +3660,8 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { for (const auto & external : match.external_bindings) { ggml::hrx::ValueBufferBinding buffer; REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = external.tensor->ne[dimension]; - nb[dimension] = external.tensor->nb[dimension]; - } external_bindings.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, external.tensor->type, ne, nb, ggml_nbytes(external.tensor) == 0)); + external.value, buffer, ggml_nbytes(external.tensor) == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); REQUIRE(bindings.valid()); @@ -3973,14 +3972,8 @@ static void run_imported_packed_weight_case(ggml_type weight_type, for (const auto & slot : match.external_bindings) { ggml::hrx::ValueBufferBinding value; REQUIRE(ggml_backend_hrx_resolve_value_buffer(slot.tensor, value)); - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = slot.tensor->ne[dimension]; - nb[dimension] = slot.tensor->nb[dimension]; - } external.push_back(ggml::hrx::make_command_program_binding( - slot.value, value, slot.tensor->type, ne, nb, ggml_nbytes(slot.tensor) == 0)); + slot.value, value, ggml_nbytes(slot.tensor) == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); REQUIRE(bindings.valid()); @@ -4078,14 +4071,8 @@ static void run_imported_packed_weight_case(ggml_type weight_type, for (const auto & slot : match.external_bindings) { ggml::hrx::ValueBufferBinding value = {}; REQUIRE(ggml_backend_hrx_resolve_value_buffer(slot.tensor, value)); - std::array ne = {}; - std::array nb = {}; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - ne[dimension] = slot.tensor->ne[dimension]; - nb[dimension] = slot.tensor->nb[dimension]; - } external.push_back(ggml::hrx::make_command_program_binding( - slot.value, value, slot.tensor->type, ne, nb, ggml_nbytes(slot.tensor) == 0)); + slot.value, value, ggml_nbytes(slot.tensor) == 0)); } return ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); }; @@ -4291,19 +4278,8 @@ static void run_imported_packed_weight_case(ggml_type weight_type, for (const auto & external : full_match.external_bindings) { ggml::hrx::ValueBufferBinding value = {}; REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, value)); - ggml::hrx::CommandProgramBinding binding = { - external.value, value.buffer, value.offset, value.length, value.identity, value.generation, - value.capacity, value.host_data, value.weight, ggml_nbytes(external.tensor) == 0, - }; - binding.source_host_data = value.host_data; - binding.source_offset = value.offset; - binding.source_capacity = value.capacity; - binding.source_type = external.tensor->type; - for (size_t dimension = 0; dimension < GGML_MAX_DIMS; ++dimension) { - binding.source_ne[dimension] = external.tensor->ne[dimension]; - binding.source_nb[dimension] = external.tensor->nb[dimension]; - } - external_bindings.push_back(binding); + external_bindings.push_back(ggml::hrx::make_command_program_binding( + external.value, value, ggml_nbytes(external.tensor) == 0)); } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); REQUIRE(bindings.valid()); From 8c7ee3f36b001a1cd478534829d3967a3ff8e9a8 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 04:32:35 +0000 Subject: [PATCH 06/18] hrx: share prepared cache entries across validated graph instances --- ggml/src/ggml-hrx/README.md | 4 +- .../dispatch/command-program-bindings.cpp | 15 -- .../dispatch/command-program-bindings.h | 5 +- ggml/src/ggml-hrx/ggml-hrx.cpp | 4 +- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 25 ++- .../ggml-hrx/runtime/graph-program-cache.cpp | 178 ++++-------------- .../ggml-hrx/runtime/graph-program-cache.h | 18 +- .../prepared-command-program-cache.cpp | 42 +++-- .../runtime/prepared-command-program-cache.h | 25 ++- tests/hrx-backend-test.cpp | 142 ++++++++++++-- tests/hrx-ops-test.cpp | 20 +- 11 files changed, 248 insertions(+), 230 deletions(-) diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index 0c7774e8f688..58d581fd5f44 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -8,7 +8,9 @@ Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx Live weight bytes count physical allocations retained by the cache, a prepared program, or a recorded graph. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. -Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. A warm graph reuses its prepared bindings without rebuilding per-weight demand; mutable inputs still rebind and upload normally. +Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. + +A graph instance remembers its validated prepared-cache entry. Its first execution resolves all physical bindings; repeated executions with the same graph UID and node array bind only mutable host staging. Structurally equivalent graphs can share command programs without sharing unchecked bindings: each new instance goes through the full binding lookup before using the fast path. Mutable source rebinding also returns to that lookup. Both execution paths use the same prepared-cache entry, lock, and replay implementation; warm executions do not rebuild per-weight demand. Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning groups exact source byte ranges; it does not detect partial overlap or infer a whole tensor from a transformed subrange. It names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 29a3193b2d06..42a27f2dcab8 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -1,6 +1,5 @@ #include "command-program-bindings.h" -#include #include namespace ggml::hrx { @@ -106,19 +105,5 @@ CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBin return { hash }; } -CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings) { - std::ostringstream out; - out << "hrx-bindings-v1"; - for (const CommandProgramBinding & binding : bindings.bindings()) { - out << "|value=" << binding.value.value << "|kind=" << (binding.host_data != nullptr ? "host" : "device") - << "|identity=" << binding.identity << "|generation=" - << (binding.tracks_content_generation() ? binding.generation : 0) - << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length - << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0) - << "|source_host=" << (binding.source_host_data != nullptr ? 1 : 0) - << "|source_offset=" << binding.source_offset << "|source_capacity=" << binding.source_capacity; - } - return { out.str() }; -} } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index 1358bb9cf9fc..bd5093a9aab9 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -43,9 +43,7 @@ struct CommandProgramBinding { CommandProgramBinding make_command_program_binding( ValueId value, const ValueBufferBinding & buffer, bool empty_value); -struct CommandProgramBindingsFingerprint { - std::string value; -}; + struct CommandProgramBindingsHash { uint64_t value = 0; @@ -71,6 +69,5 @@ class CommandProgramBindings { }; CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBindings & bindings); -CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const CommandProgramBindings & bindings); } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index f448c1ba2194..68f3c61ed4d7 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -1333,8 +1333,8 @@ bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_c const ggml::hrx::HostTransferStats transfer_stats = context->host_transfers.stats(); stats->graph_program_builds = graph_stats.builds; stats->graph_program_hits = graph_stats.hits; - stats->prepared_program_builds = graph_stats.prepared_program_builds + prepared_stats.builds; - stats->prepared_program_hits = graph_stats.prepared_program_hits + prepared_stats.hits; + stats->prepared_program_builds = prepared_stats.builds; + stats->prepared_program_hits = prepared_stats.hits; stats->weight_materializations = weight_stats.misses; stats->weight_reuses = weight_stats.hits; stats->weight_uploaded_bytes = weight_stats.materialized_bytes; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 044010b5e011..4a49629ebb84 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -85,9 +85,7 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { return result; } - const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); - const bool has_prepared = lookup.program->has_prepared_program(); - const bool reusable_owner = use_graph_prepared && has_prepared; + const bool reusable_owner = lookup.program->can_use_prepared_fast_path(graph); const bool mutable_sources_current = reusable_owner && lookup.program->mutable_sources_current(); const bool trusted_weights = reusable_owner && mutable_sources_current; GraphProgramMatch binding_match = std::move(lookup.match); @@ -117,11 +115,17 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { &context_.device->host_buffers, &context_.graph_replay_state, }; - const PreparedCommandProgramCacheExecutionResult execution = - use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings, !trusted_weights) : - context_.prepared_programs.execute_with_result(execution_context, lookup.program->uid(), - lookup.program->command_shape_hash(), - lookup.program->commands(), bindings); + std::shared_ptr entry; + if (trusted_weights) { + entry = lookup.program->prepared_entry(); + } else if (graph.uid != 0) { + entry = context_.prepared_programs.get_or_create(execution_context, lookup.program->uid(), + lookup.program->command_shape_hash(), bindings); + } + const PreparedCommandProgramCacheExecutionResult execution = entry != nullptr ? + context_.prepared_programs.execute_with_result(execution_context, lookup.program->commands(), bindings, entry) : + context_.prepared_programs.execute_with_result(execution_context, 0, lookup.program->command_shape_hash(), + lookup.program->commands(), bindings); result.graph_replay_event = execution.graph_replay_event; if (execution.execution_resources != nullptr) { // A failed return can still follow queued writes. @@ -135,6 +139,11 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { return result; } + const bool validated_entry = entry != nullptr && !trusted_weights; + if (validated_entry) { + lookup.program->bind_prepared_graph(graph, binding_match, std::move(entry)); + } + result.code = GGML_STATUS_SUCCESS; return result; } diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index 49fafd29de96..53d9d2037da4 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -5,8 +5,6 @@ #include "dispatch/dispatch-scheduler.h" #include "ggml-impl.h" #include "ggml.h" -#include "hrx-interop-utils.h" -#include "hrx_runtime.h" #include #include @@ -162,21 +160,6 @@ static bool environment_flag_enabled(const char * name) { return value != nullptr && value[0] != '\0' && value[0] != '0'; } -static void apply_graph_replay_result(PreparedCommandProgramCacheExecutionResult & result, - const RecordedCommandGraphExecutionResult & replay) { - result.graph_replay_event = replay.event; - result.graph_replay_ineligible_reason = replay.ineligible_reason; - result.graph_replay_build_ns = replay.build_ns; - result.graph_replay_launch_ns = replay.launch_ns; - result.graph_replay_total_ns = replay.total_ns(); - result.graph_replay_dispatches = replay.dispatch_count; - result.graph_replay_transient_allocation_changed = replay.transient_allocation_changed; -} - -static bool graph_replay_should_fallback(HrxGraphReplayEvent event) { - return event == HrxGraphReplayEvent::Ineligible || event == HrxGraphReplayEvent::BuildFailed; -} - static void collect_command_graph_values(const std::vector & commands, std::vector & values) { for (const Command & command : commands) { for (const CommandBinding & binding : command.bindings) { @@ -424,12 +407,14 @@ GraphProgramMatch GraphProgram::match_trusted_graph(const ggml_cgraph & current_ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & current_graph) const { GraphProgramMatch result; std::lock_guard lock(prepared_mutex_); - if (!has_prepared_) { + if (prepared_entry_ == nullptr) { result.status.log("missing prepared HRX command program"); return result; } - result.external_bindings.reserve(resources_->prepared.host_staging.size()); - for (const HostStagingBuffer & staging : resources_->prepared.host_staging) { + std::lock_guard entry_lock(prepared_entry_->mutex); + const auto & staging_buffers = prepared_entry_->resources->prepared.host_staging; + result.external_bindings.reserve(staging_buffers.size()); + for (const HostStagingBuffer & staging : staging_buffers) { const GraphProgramExternalSlot * slot = find_external_slot(ValueId(staging.value)); if (slot == nullptr) { result.status.log("prepared host staging value %d has no external graph slot", staging.value); @@ -447,11 +432,9 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { Status status; external_slots_.clear(); - mutable_source_watches_.clear(); written_slots_.clear(); std::unordered_set written_values; external_slot_by_value_.clear(); - fast_path_nodes_ = graph.nodes; external_slots_.reserve(match.external_bindings.size()); auto collect_access = [&](const std::vector & commands) { for (const Command & command : commands) { @@ -470,16 +453,6 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra collect_access(commands_->initialization_commands); collect_access(commands_->commands); for (const GraphProgramExternalBinding & binding : match.external_bindings) { - ValueBufferBinding value_binding; - const bool bound = ggml_backend_hrx_resolve_value_buffer(binding.tensor, value_binding); - const bool weight = bound && value_binding.weight; - ggml_backend_hrx_imported_region region = {}; - const bool mutable_source = bound && !weight; - const bool retained_source = mutable_source && ggml_backend_hrx_get_retained_region(binding.tensor, region); - if (retained_source) { - // Mutable inputs may rebind under one graph UID; immutable weight bindings are trusted. - mutable_source_watches_.push_back({ binding.tensor, region.identity, region.offset, region.length }); - } GraphProgramExternalSlot slot; slot.value = binding.value; slot.weight_layout_eligible = ggml_backend_hrx_weight_layout_eligible(binding.tensor); @@ -521,15 +494,38 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra return status; } -bool GraphProgram::has_prepared_program() const { +std::shared_ptr GraphProgram::prepared_entry() const { + std::lock_guard lock(prepared_mutex_); + return prepared_entry_; +} + +void GraphProgram::bind_prepared_graph(const ggml_cgraph & graph, const GraphProgramMatch & match, + std::shared_ptr entry) { std::lock_guard lock(prepared_mutex_); - return has_prepared_; + mutable_source_watches_.clear(); + for (const GraphProgramExternalBinding & binding : match.external_bindings) { + ValueBufferBinding value_binding; + const bool bound = ggml_backend_hrx_resolve_value_buffer(binding.tensor, value_binding); + const bool weight = bound && value_binding.weight; + ggml_backend_hrx_imported_region region = {}; + const bool mutable_source = bound && !weight; + const bool retained_source = mutable_source && ggml_backend_hrx_get_retained_region(binding.tensor, region); + if (retained_source) { + // Mutable inputs may rebind under one graph UID; immutable weight bindings are trusted. + mutable_source_watches_.push_back({ binding.tensor, region.identity, region.offset, region.length }); + } + } + // The full binding lookup established this graph instance's physical resources. + prepared_entry_ = std::move(entry); + fast_path_uid_ = graph.uid; + fast_path_nodes_ = graph.nodes; } bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { + std::lock_guard lock(prepared_mutex_); // The UID prevents graph-arena address reuse from passing the fast path. const bool allocated_graph = graph.uid != 0; - const bool same_uid = graph.uid == uid_; + const bool same_uid = graph.uid == fast_path_uid_; const bool same_nodes = graph.nodes == fast_path_nodes_; return allocated_graph && same_uid && same_nodes; } @@ -574,108 +570,6 @@ void GraphProgram::mark_imported_writes(const ggml_cgraph & graph) const { } } -PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { - std::lock_guard lock(prepared_mutex_); - return prepared_stats_; -} - -PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( - const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings, - bool full_bindings) { - PreparedCommandProgramCacheExecutionResult result; - std::lock_guard lock(prepared_mutex_); - if (commands_ == nullptr || !commands_->valid() || !bindings.valid()) { - result.graph_replay_event = HrxGraphReplayEvent::Ineligible; - result.graph_replay_ineligible_reason = "invalid_graph_program"; - if (commands_ == nullptr) { - result.status.log("missing cached HRX command program"); - } - result.status.append(bindings.status); - return result; - } - - const std::string bindings_fingerprint = - full_bindings ? command_program_bindings_fingerprint(bindings).value : std::string{}; - const bool fingerprint_changed = prepared_bindings_fingerprint_ != bindings_fingerprint; - const bool prepared_bindings_changed = full_bindings && fingerprint_changed; - const bool replace_prepared = has_prepared_ && prepared_bindings_changed; - if (replace_prepared) { - // A retained source may have changed while the graph structure stayed identical. Retire old bindings only - // after queued work completes; replay captures physical addresses that cannot be rebound in place. - if (ErrorResult error = take_status(hrx_stream_synchronize(context.stream))) { - result.status.log("synchronize before replacing HRX weight bindings failed: %s", error->c_str()); - return result; - } - context.graph_replay_state->mark_stream_synchronized(); - resources_.reset(); - has_prepared_ = false; - } - if (!has_prepared_) { - resources_ = std::make_shared(); - resources_->prepared = prepare_command_program(context, *commands_, bindings); - if (!resources_->prepared.valid()) { - result.status.append(resources_->prepared.status); - resources_.reset(); - return result; - } - has_prepared_ = true; - prepared_bindings_fingerprint_ = bindings_fingerprint; - for (MutableSourceWatch & watch : mutable_source_watches_) { - ggml_backend_hrx_imported_region region = {}; - if (ggml_backend_hrx_get_retained_region(watch.tensor, region)) { - watch.identity = region.identity; - watch.offset = region.offset; - watch.length = region.length; - } - } - ++prepared_stats_.builds; - } else { - ++prepared_stats_.hits; - } - result.execution_resources = resources_; - - const bool replay_disabled = graph_replay_disabled_by_environment(); - const bool serial_execution = debug_serial_command_execution_enabled(); - const bool execute_directly = replay_disabled || serial_execution; - if (execute_directly) { - result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = serial_execution ? - "debug_serial_execution" : - "disabled_by_environment"; - result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, resources_->prepared); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } - return result; - } - - if (environment_flag_enabled("GGML_HRX_REBUILD_GRAPH_REPLAY")) { - resources_->recorded = {}; - } - - const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, *commands_, bindings, resources_->prepared, - resources_->recorded); - apply_graph_replay_result(result, replay); - if (replay.success) { - result.success = true; - return result; - } - if (!graph_replay_should_fallback(replay.event)) { - result.status.append(replay.status); - if (result.status.success()) { - result.status.log("execute cached HRX graph replay failed"); - } - return result; - } - result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, resources_->prepared); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } - return result; -} - GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_graph) const { GraphProgramMatch result; if (graph_ == nullptr) { @@ -881,7 +775,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, } if (cached_program != nullptr) { const bool bind_external = - !cached_program->has_prepared_program() || !cached_program->can_use_prepared_fast_path(graph); + !cached_program->can_use_prepared_fast_path(graph); GraphProgramMatch match = disable_fast_path || validate_fast_path ? cached_program->match_current_graph(graph) : cached_program->match_trusted_graph(graph, bind_external); @@ -946,13 +840,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, GraphProgramCacheStats GraphProgramCache::stats() const { std::lock_guard lock(mutex_); - GraphProgramCacheStats stats = stats_; - for (const auto & entry : programs_) { - const PreparedCommandProgramCacheStats prepared = entry.second->prepared_stats(); - stats.prepared_program_builds += prepared.builds; - stats.prepared_program_hits += prepared.hits; - } - return stats; + return stats_; } void GraphProgramCache::clear() { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index e67a348eaec2..d59c6e1a7aee 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -73,16 +73,12 @@ class GraphProgram { GraphProgramMatch match_host_staging_graph(const ggml_cgraph & graph) const; Status capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match); - bool has_prepared_program() const; bool can_use_prepared_fast_path(const ggml_cgraph & graph) const; bool mutable_sources_current() const; void mark_imported_writes(const ggml_cgraph & graph) const; - - PreparedCommandProgramCacheStats prepared_stats() const; - - PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings, - bool full_bindings); + std::shared_ptr prepared_entry() const; + void bind_prepared_graph(const ggml_cgraph & graph, const GraphProgramMatch & match, + std::shared_ptr entry); private: const GraphProgramExternalSlot * find_external_slot(ValueId value) const; @@ -107,19 +103,15 @@ class GraphProgram { std::vector mutable_source_watches_; std::vector written_slots_; std::unordered_map external_slot_by_value_; + uint64_t fast_path_uid_ = 0; const ggml_tensor * const * fast_path_nodes_ = nullptr; mutable std::mutex prepared_mutex_; - std::shared_ptr resources_; - std::string prepared_bindings_fingerprint_; - PreparedCommandProgramCacheStats prepared_stats_; - bool has_prepared_ = false; + std::shared_ptr prepared_entry_; }; struct GraphProgramCacheStats { uint64_t builds = 0; uint64_t hits = 0; - uint64_t prepared_program_builds = 0; - uint64_t prepared_program_hits = 0; }; struct GraphProgramLookup { diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index 817daa15a477..31fb48bdeb54 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -114,19 +115,32 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ return result; } - const Key key = cache_key(graph_uid, context, command_shape_hash, bindings); - std::shared_ptr entry; - { - std::lock_guard lock(mutex_); - auto found = programs_.find(key); - if (found == programs_.end()) { - entry = std::make_shared(); - programs_.emplace(key, entry); - } else { - entry = found->second; - } + return execute_with_result(context, commands, bindings, + get_or_create(context, graph_uid, command_shape_hash, bindings)); +} + +std::shared_ptr PreparedCommandProgramCache::get_or_create( + const CommandProgramExecutionContext & context, + uint64_t graph_uid, + uint64_t command_shape_hash, + const CommandProgramBindings & bindings) { + const Key key = cache_key(graph_uid, context, command_shape_hash, bindings); + std::lock_guard lock(mutex_); + auto found = programs_.find(key); + if (found != programs_.end()) { + return found->second; } + auto entry = std::make_shared(); + programs_.emplace(key, entry); + return entry; +} +PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_with_result( + const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + const std::shared_ptr & entry) { + PreparedCommandProgramCacheExecutionResult result; std::lock_guard entry_lock(entry->mutex); const bool replay_disabled = graph_replay_disabled_by_environment(); const bool serial_execution = debug_serial_command_execution_enabled(); @@ -145,6 +159,12 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ } return result; } + const char * rebuild_setting = std::getenv("GGML_HRX_REBUILD_GRAPH_REPLAY"); + const bool rebuild_requested = rebuild_setting != nullptr && rebuild_setting[0] != '\0'; + const bool rebuild_enabled = rebuild_requested && rebuild_setting[0] != '0'; + if (rebuild_enabled) { + resources.recorded = {}; + } const RecordedCommandGraphExecutionResult replay = bind_and_launch_recorded_command_graph(context, commands, bindings, resources.prepared, resources.recorded); apply_graph_replay_result(result, replay); diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index 5a0a5c2fef37..a5f713301d8a 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -22,6 +22,11 @@ struct PreparedCommandProgramCacheStats { uint64_t hits = 0; }; +struct PreparedCommandProgramCacheEntry { + std::mutex mutex; + std::shared_ptr resources; +}; + struct PreparedCommandProgramCacheExecutionResult { bool success = false; Status status; @@ -37,6 +42,19 @@ struct PreparedCommandProgramCacheExecutionResult { class PreparedCommandProgramCache { public: + // Full physical bindings select the entry; a validated graph instance can retain it for warm execution. + std::shared_ptr get_or_create( + const CommandProgramExecutionContext & context, + uint64_t graph_uid, + uint64_t command_shape_hash, + const CommandProgramBindings & bindings); + + PreparedCommandProgramCacheExecutionResult execute_with_result( + const CommandProgramExecutionContext & context, + const CommandProgram & commands, + const CommandProgramBindings & bindings, + const std::shared_ptr & entry); + bool execute(const CommandProgramExecutionContext & context, uint64_t graph_uid, const std::string & command_shape, @@ -81,16 +99,11 @@ class PreparedCommandProgramCache { uint64_t command_shape_hash, const CommandProgramBindings & bindings) const; - struct Entry { - std::mutex mutex; - std::shared_ptr resources; - }; - void record_build(); void record_hit(); mutable std::mutex mutex_; - std::unordered_map, KeyHash> programs_; + std::unordered_map, KeyHash> programs_; PreparedCommandProgramCacheStats stats_; }; diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 65e399e0eb0f..4247d19489d9 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1015,9 +1015,7 @@ static void run_graph_import_checks() { REQUIRE(b_binding->length == b_value->byte_count); REQUIRE(runtime_bindings.find(ggml::hrx::ValueId(123456)) == nullptr); - const ggml::hrx::CommandProgramBindingsFingerprint runtime_fingerprint = - ggml::hrx::command_program_bindings_fingerprint(runtime_bindings); - REQUIRE(!runtime_fingerprint.value.empty()); + const auto runtime_hash = ggml::hrx::command_program_bindings_hash(runtime_bindings); ggml::hrx::ValueMap changed_identity = imported.graph.values(); REQUIRE(changed_identity.bind_buffer( @@ -1025,8 +1023,8 @@ static void run_graph_import_checks() { const ggml::hrx::CommandProgramBindings changed_identity_bindings = ggml::hrx::CommandProgramBindings::from_value_map(changed_identity); REQUIRE(changed_identity_bindings.valid()); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_identity_bindings).value != - runtime_fingerprint.value); + REQUIRE(ggml::hrx::command_program_bindings_hash(changed_identity_bindings).value != + runtime_hash.value); REQUIRE(!a_binding->weight); std::vector changed_generation_values = runtime_bindings.bindings(); @@ -1038,9 +1036,7 @@ static void run_graph_import_checks() { ggml::hrx::CommandProgramBindings::from_bindings(std::move(changed_generation_values)); REQUIRE(changed_generation_bindings.valid()); REQUIRE(ggml::hrx::command_program_bindings_hash(changed_generation_bindings).value == - ggml::hrx::command_program_bindings_hash(runtime_bindings).value); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_generation_bindings).value == - runtime_fingerprint.value); + runtime_hash.value); ggml::hrx::ValueMap changed_capacity = imported.graph.values(); REQUIRE(changed_capacity.bind_buffer( @@ -1048,8 +1044,8 @@ static void run_graph_import_checks() { const ggml::hrx::CommandProgramBindings changed_capacity_bindings = ggml::hrx::CommandProgramBindings::from_value_map(changed_capacity); REQUIRE(changed_capacity_bindings.valid()); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_capacity_bindings).value != - runtime_fingerprint.value); + REQUIRE(ggml::hrx::command_program_bindings_hash(changed_capacity_bindings).value != + runtime_hash.value); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -1182,7 +1178,7 @@ static void run_graph_import_checks() { REQUIRE(staging_binding != nullptr); REQUIRE(staging_binding->host_data == host_storage.data()); REQUIRE(staging_binding->source_host_data == nullptr); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(host_bindings).value != runtime_fingerprint.value); + REQUIRE(ggml::hrx::command_program_bindings_hash(host_bindings).value != runtime_hash.value); std::vector changed_weight_values = host_bindings.bindings(); auto changed_weight = std::find_if(changed_weight_values.begin(), changed_weight_values.end(), [&](const auto & binding) { return binding.value == a_value->id; }); @@ -1193,8 +1189,6 @@ static void run_graph_import_checks() { REQUIRE(changed_weight_bindings.valid()); REQUIRE(ggml::hrx::command_program_bindings_hash(changed_weight_bindings).value != ggml::hrx::command_program_bindings_hash(host_bindings).value); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(changed_weight_bindings).value != - ggml::hrx::command_program_bindings_fingerprint(host_bindings).value); resolved = ggml::hrx::resolve_command_program_bindings(commands, missing_bindings); REQUIRE(!resolved.valid()); @@ -11769,7 +11763,9 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); std::vector expected(width); ggml_backend_tensor_get(cpu_output, expected.data(), 0, ggml_nbytes(cpu_output)); - auto execution = lookup.program->execute_with_result(execution_context, bindings, true); + auto execution = state->prepared_programs.execute_with_result( + execution_context, lookup.program->uid(), lookup.program->command_shape_hash(), + lookup.program->commands(), bindings); REQUIRE(execution.success); ggml_backend_synchronize(backend); const auto & prepared = execution.execution_resources->prepared; @@ -11783,8 +11779,8 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { // The vector matmul declares both host input and output as reads. const size_t staging_upload_bytes = ggml_nbytes(input) + ggml_nbytes(output); REQUIRE(state->host_transfers.stats().upload_bytes == weight_bytes + (iteration + 1) * staging_upload_bytes); - REQUIRE(lookup.program->prepared_stats().builds == 1); - REQUIRE(lookup.program->prepared_stats().hits == static_cast(iteration)); + REQUIRE(state->prepared_programs.stats().builds == 1); + REQUIRE(state->prepared_programs.stats().hits == static_cast(iteration)); const auto expected_event = direct_execution ? ggml::hrx::HrxGraphReplayEvent::Disabled : (iteration == 0 ? ggml::hrx::HrxGraphReplayEvent::MissBuild : ggml::hrx::HrxGraphReplayEvent::Hit); REQUIRE(execution.graph_replay_event == expected_event); @@ -12957,6 +12953,119 @@ static void run_validated_graph_uid_match_checks() { ggml_free(ctx); } +static void run_prepared_graph_instance_reuse() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * state = static_cast(backend->context); + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + struct TestGraph { + ggml_context * context = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_tensor * left = nullptr; + ggml_tensor * right = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + std::array storage = {}; + ~TestGraph() { + ggml_backend_buffer_free(buffer); + ggml_free(context); + } + }; + auto make_graph = [&](uint64_t uid) { + auto item = std::make_unique(); + item->context = ggml_init({ 256 * 1024, nullptr, true }); + REQUIRE(item->context != nullptr); + item->left = ggml_new_tensor_1d(item->context, GGML_TYPE_F32, 8); + item->right = ggml_new_tensor_1d(item->context, GGML_TYPE_F32, 8); + item->output = ggml_add(item->context, item->left, item->right); + ggml_set_output(item->output); + item->graph = ggml_new_graph(item->context); + ggml_build_forward_expand(item->graph, item->output); + item->graph->uid = uid; + item->buffer = ggml_backend_dev_buffer_from_host_ptr(ggml_backend_get_device(backend), + item->storage.data(), sizeof(item->storage), 8 * sizeof(float)); + REQUIRE(item->buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(item->buffer, item->left, item->storage.data()) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(item->buffer, item->right, item->storage.data() + 8) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_tensor_alloc(item->buffer, item->output, item->storage.data() + 16) == GGML_STATUS_SUCCESS); + return item; + }; + auto compute = [&](TestGraph & item, float seed) { + std::array left, right, expected, actual; + for (size_t i = 0; i < left.size(); ++i) { + left[i] = seed + i; + right[i] = seed * 2 - i; + expected[i] = left[i] + right[i]; + } + ggml_backend_tensor_set(item.left, left.data(), 0, sizeof(left)); + ggml_backend_tensor_set(item.right, right.data(), 0, sizeof(right)); + REQUIRE(ggml_backend_graph_compute(backend, item.graph) == GGML_STATUS_SUCCESS); + ggml_backend_tensor_get(item.output, actual.data(), 0, sizeof(actual)); + REQUIRE(actual == expected); + }; + auto first = make_graph(4301); + auto second = make_graph(4302); + auto lookup = state->graph_programs.get_or_build(*first->graph, corpus, state->device->architecture); + REQUIRE(lookup.valid()); + auto * program = lookup.program; + REQUIRE(!program->can_use_prepared_fast_path(*first->graph)); + compute(*first, 1); + REQUIRE(program->can_use_prepared_fast_path(*first->graph)); + auto first_entry = program->prepared_entry(); + compute(*first, 2); + REQUIRE(state->prepared_programs.stats().builds == 1); + REQUIRE(state->prepared_programs.stats().hits == 1); + + compute(*second, 3); + REQUIRE(state->graph_programs.get_or_build(*second->graph, corpus, state->device->architecture).program == program); + REQUIRE(program->uid() == 4301); + REQUIRE(program->can_use_prepared_fast_path(*second->graph)); + REQUIRE(!program->can_use_prepared_fast_path(*first->graph)); + auto second_entry = program->prepared_entry(); + REQUIRE(first_entry != second_entry); + compute(*second, 4); + REQUIRE(state->prepared_programs.stats().builds == 2); + compute(*first, 5); + REQUIRE(program->prepared_entry() == first_entry); + compute(*second, 6); + REQUIRE(program->prepared_entry() == second_entry); + REQUIRE(state->prepared_programs.stats().builds == 2); + + // The current instance must own the watches even after the original graph arena is gone. + first.reset(); + second->left->data = second->storage.data() + 24; + REQUIRE(!program->mutable_sources_current()); + compute(*second, 7); + REQUIRE(program->mutable_sources_current()); + REQUIRE(state->prepared_programs.stats().builds == 3); + compute(*second, 8); + REQUIRE(state->prepared_programs.stats().builds == 3); + + std::array replacement = {}; + ggml_backend_buffer_t rebound = ggml_backend_dev_buffer_from_host_ptr(ggml_backend_get_device(backend), + replacement.data(), sizeof(replacement), sizeof(replacement)); + REQUIRE(rebound != nullptr); + second->left->buffer = rebound; + second->left->data = replacement.data(); + REQUIRE(!program->mutable_sources_current()); + compute(*second, 9); + REQUIRE(program->mutable_sources_current()); + REQUIRE(state->prepared_programs.stats().builds == 4); + + // Reusing the node-array address with a new UID must require a new binding check. + second->graph->uid = 4303; + REQUIRE(!program->can_use_prepared_fast_path(*second->graph)); + compute(*second, 10); + REQUIRE(program->can_use_prepared_fast_path(*second->graph)); + REQUIRE(state->prepared_programs.stats().builds == 4); + REQUIRE(state->graph_programs.stats().builds == 1); + first_entry.reset(); + second_entry.reset(); + second.reset(); + ggml_backend_buffer_free(rebound); + ggml_backend_free(backend); +} + static void run_graph_executor_contract_checks() { ggml_backend_hrx_device_context device_context = {}; ggml_backend_hrx_context backend_context = {}; @@ -14706,6 +14815,7 @@ static void register_hrx_backend_device_cases(test_runner::Suite & suite) { suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); suite.device_case("zero_output_unused_weight_staging.recorded", [] { run_zero_output_unused_weight_staging(false); }); suite.device_case("zero_output_unused_weight_staging.direct", [] { run_zero_output_unused_weight_staging(true); }); + suite.device_case("prepared_graph_instance_reuse", [] { run_prepared_graph_instance_reuse(); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); suite.device_case("binary_q8_publication_numerics", [] { run_binary_q8_publication_numerics(); }); suite.device_case("qwen_expert_table_partition_prefill_512_execution", diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 3c50fbcc36c9..b7f5e918be95 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3376,7 +3376,8 @@ static void run_weight_persistence_overlap_case() { &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, }; - auto execution = lookup.program->execute_with_result(context, bindings, true); + auto execution = state->prepared_programs.execute_with_result( + context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), bindings); REQUIRE(execution.success); REQUIRE(execution.execution_resources->recorded.valid()); REQUIRE(execution.execution_resources->recorded.resident_host_weights == @@ -3696,10 +3697,9 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, }; - auto execution = uncached ? - state->prepared_programs.execute_with_result(execution_context, 0, packed_lookup.program->command_shape_hash(), - packed_lookup.program->commands(), packed_bindings) : - packed_lookup.program->execute_with_result(execution_context, packed_bindings, true); + auto execution = state->prepared_programs.execute_with_result( + execution_context, uncached ? 0 : packed_lookup.program->uid(), packed_lookup.program->command_shape_hash(), + packed_lookup.program->commands(), packed_bindings); REQUIRE(execution.success); REQUIRE(!execution.execution_resources->prepared.host_staging.empty()); REQUIRE(execution.execution_resources->recorded.valid() == !uncached); @@ -3756,6 +3756,7 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { REQUIRE(completed.allocation_count == 1); REQUIRE(completed.resident_bytes == packed_source.materialized_length); state->graph_programs.clear(); + state->prepared_programs.clear(); REQUIRE(sole_resource.expired()); REQUIRE(state->host_weights.stats().allocation_count == 1); state->host_weights.clear(); @@ -4088,7 +4089,8 @@ static void run_imported_packed_weight_case(ggml_type weight_type, auto staging_bindings = bindings_for_match(staging_match); REQUIRE(staging_bindings.valid()); { - const auto graph_replay = lookup.program->execute_with_result(replay_context, staging_bindings, false); + const auto graph_replay = backend_state->prepared_programs.execute_with_result( + replay_context, lookup.program->commands(), staging_bindings, lookup.program->prepared_entry()); REQUIRE(graph_replay.success); const char * replay_setting = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); const bool replay_disabled = replay_setting != nullptr && replay_setting[0] != '\0' && replay_setting[0] != '0'; @@ -4446,9 +4448,9 @@ static void run_mixed_weight_layout_program_case(ggml_type weight_type) { const auto repeated = state->host_weights.stats(); REQUIRE(repeated.misses == first.misses); REQUIRE(repeated.duplicate_warnings == first.duplicate_warnings); - const auto repeated_graph_stats = state->graph_programs.stats(); - REQUIRE(repeated_graph_stats.prepared_program_hits > 0); - REQUIRE(repeated_graph_stats.prepared_program_builds >= 1); + const auto repeated_prepared_stats = state->prepared_programs.stats(); + REQUIRE(repeated_prepared_stats.hits > 0); + REQUIRE(repeated_prepared_stats.builds >= 1); ggml_backend_hrx_cache_stats public_repeated = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &public_repeated)); REQUIRE(public_repeated.weight_duplicate_warnings == public_first.weight_duplicate_warnings); From 827a95c8602481920ff1b64569b662fdc54363aa Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 04:41:54 +0000 Subject: [PATCH 07/18] test: account for shared prepared entries in overlap coverage --- tests/hrx-ops-test.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 12a0525cb9f2..6ed99c6fbd0d 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3441,7 +3441,9 @@ static void run_weight_persistence_overlap_case() { REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); REQUIRE(warmed.weight_live_allocations == 2); REQUIRE(warmed.weight_duplicate_warnings >= cross_queued.weight_duplicate_warnings); - REQUIRE(state->prepared_programs.stats().builds == 0); + // Graph execution and the direct submission helper share the warmed native entry. + REQUIRE(state->prepared_programs.stats().builds == 3); + REQUIRE(state->prepared_programs.stats().hits >= 1); // A cold preparation failure must leave an already queued native reader's lease intact. REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); From 413303d1b862487124a3087312fa7c911f793da8 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 06:00:41 +0000 Subject: [PATCH 08/18] hrx: capture mutable watches from resolved bindings --- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 2 +- .../ggml-hrx/runtime/graph-program-cache.cpp | 17 +++++++++-------- ggml/src/ggml-hrx/runtime/graph-program-cache.h | 1 + 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index f87c51301534..3673f1c7c11b 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -144,7 +144,7 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { const bool validated_entry = entry != nullptr && !trusted_weights; if (validated_entry) { - lookup.program->bind_prepared_graph(graph, binding_match, std::move(entry)); + lookup.program->bind_prepared_graph(graph, binding_match, bindings, std::move(entry)); } result.code = GGML_STATUS_SUCCESS; diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index 53d9d2037da4..2c83f7d94998 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -500,19 +500,20 @@ std::shared_ptr GraphProgram::prepared_entry() } void GraphProgram::bind_prepared_graph(const ggml_cgraph & graph, const GraphProgramMatch & match, + const CommandProgramBindings & bindings, std::shared_ptr entry) { std::lock_guard lock(prepared_mutex_); mutable_source_watches_.clear(); - for (const GraphProgramExternalBinding & binding : match.external_bindings) { - ValueBufferBinding value_binding; - const bool bound = ggml_backend_hrx_resolve_value_buffer(binding.tensor, value_binding); - const bool weight = bound && value_binding.weight; - ggml_backend_hrx_imported_region region = {}; - const bool mutable_source = bound && !weight; - const bool retained_source = mutable_source && ggml_backend_hrx_get_retained_region(binding.tensor, region); + // Successful external binding preserves match order and already resolved every physical source. + for (size_t i = 0; i < match.external_bindings.size(); ++i) { + const CommandProgramBinding & binding = bindings.bindings()[i]; + const bool mutable_source = !binding.weight; + const bool host_source = binding.source_host_data != nullptr; + const bool retained_source = mutable_source && host_source; if (retained_source) { // Mutable inputs may rebind under one graph UID; immutable weight bindings are trusted. - mutable_source_watches_.push_back({ binding.tensor, region.identity, region.offset, region.length }); + mutable_source_watches_.push_back({ match.external_bindings[i].tensor, binding.identity, + binding.source_offset, binding.length }); } } // The full binding lookup established this graph instance's physical resources. diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index d59c6e1a7aee..d6d57945b6dd 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -78,6 +78,7 @@ class GraphProgram { void mark_imported_writes(const ggml_cgraph & graph) const; std::shared_ptr prepared_entry() const; void bind_prepared_graph(const ggml_cgraph & graph, const GraphProgramMatch & match, + const CommandProgramBindings & bindings, std::shared_ptr entry); private: From 463fd0b285835d36cf15544b7bc121fc37400ca7 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 06:14:58 +0000 Subject: [PATCH 09/18] hrx: inline physical command binding construction --- .../ggml-hrx/dispatch/command-program-bindings.cpp | 11 ----------- .../src/ggml-hrx/dispatch/command-program-bindings.h | 12 ++++++++++-- 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 42a27f2dcab8..a5c01da21925 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -12,17 +12,6 @@ static void mix_hash(uint64_t & hash, uint64_t value) { } // namespace -CommandProgramBinding make_command_program_binding( - ValueId value, const ValueBufferBinding & buffer, bool empty_value) { - CommandProgramBinding binding = { value, buffer.buffer, buffer.offset, buffer.length, buffer.identity, - buffer.generation, buffer.capacity, buffer.host_data, buffer.weight, - empty_value }; - binding.source_host_data = buffer.source_host_data; - binding.source_offset = buffer.source_offset; - binding.source_capacity = buffer.source_capacity; - return binding; -} - CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & values) { std::vector bindings; CommandProgramBindings result; diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index bd5093a9aab9..ddf650ecec99 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -40,8 +40,16 @@ struct CommandProgramBinding { bool tracks_content_generation() const { return weight && source_host_data != nullptr; } }; -CommandProgramBinding make_command_program_binding( - ValueId value, const ValueBufferBinding & buffer, bool empty_value); +inline CommandProgramBinding make_command_program_binding( + ValueId value, const ValueBufferBinding & buffer, bool empty_value) { + CommandProgramBinding binding = { value, buffer.buffer, buffer.offset, buffer.length, buffer.identity, + buffer.generation, buffer.capacity, buffer.host_data, buffer.weight, + empty_value }; + binding.source_host_data = buffer.source_host_data; + binding.source_offset = buffer.source_offset; + binding.source_capacity = buffer.source_capacity; + return binding; +} From 16a4233049fad94b6f159697cab6eeebc6c1a6e1 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 21:41:00 +0000 Subject: [PATCH 10/18] tests: remove HRX-specific architecture test Assisted-by: Codex --- tests/test-llama-archs.cpp | 110 +------------------------------------ 1 file changed, 1 insertion(+), 109 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index f2f4a05e6e12..aff6a7c90b68 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -7,10 +7,6 @@ #include "llama.h" #include "llama-cpp.h" #include "../src/llama-ext.h" -#ifdef GGML_USE_HRX -#include "ggml-hrx.h" -#include "../src/llama-model.h" -#endif #include #include @@ -160,7 +156,6 @@ static void usage(char ** argv) { LOG(" -o, --out Save generated test models to instead of running backend tests\n"); LOG(" -v Set log verbosity level\n"); LOG(" -b, --backend Run only on the given backend device\n"); - LOG(" --hrx-import Test mmap-backed HRX weight import\n"); LOG(" -h, --help Show this help message\n\n"); LOG("Examples:\n"); LOG(" %s\n", argv[0]); @@ -231,7 +226,7 @@ struct kv_writer { } }; -static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool hrx_compatible = false) { +static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { gguf_context_ptr ret(gguf_init_empty()); kv_writer ms(arch, ret.get()); const uint32_t n_ctx = 256; @@ -241,9 +236,6 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const uint32_t n_head = 2; uint32_t n_ff = 384; uint32_t n_layer = 2; - if (hrx_compatible) { - n_ff = 256; - } if (arch == LLM_ARCH_LLAMA4) { n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4 } else if (arch == LLM_ARCH_MAION_CODER) { @@ -1376,100 +1368,6 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con return all_ok ? 0 : 1; } -static int test_hrx_mmap_import(const size_t seed, const float stdev) { -#ifndef GGML_USE_HRX - GGML_UNUSED(seed); - GGML_UNUSED(stdev); - throw std::runtime_error("mmap import test requires an HRX build"); -#else - ggml_backend_load_all(); - ggml_backend_dev_t hrx_device = nullptr; - for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { - ggml_backend_dev_t device = ggml_backend_dev_get(i); - if (std::strncmp(ggml_backend_dev_name(device), "HRX", 3) == 0) { - hrx_device = device; - break; - } - } - if (hrx_device == nullptr) { - throw std::runtime_error("HRX device unavailable for mmap import test"); - } - ggml_backend_dev_props props = {}; - ggml_backend_dev_get_props(hrx_device, &props); - if (!props.caps.buffer_from_host_ptr) { - throw std::runtime_error("HRX host-pointer import is not advertised"); - } - - gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_LLAMA, false, true); - auto cpu_model = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, {}); - const ggml_tensor * cpu_weight = cpu_model.first->get_tensor("blk.0.attn_q.weight"); - if (cpu_weight == nullptr) { - throw std::runtime_error("tiny GGUF has no attention weight"); - } - std::vector expected(ggml_nbytes(cpu_weight)); - ggml_backend_tensor_get(cpu_weight, expected.data(), 0, expected.size()); - - const std::filesystem::path path = std::filesystem::temp_directory_path() / - ("hrx-import-" + std::to_string(seed) + ".gguf"); - struct RemoveFile { - std::filesystem::path path; - ~RemoveFile() { std::error_code ignored; std::filesystem::remove(path, ignored); } - } remove_file{ path }; - llama_model_save_to_file(cpu_model.first.get(), path.string().c_str()); - - ggml_backend_dev_t devices[] = { hrx_device, nullptr }; - llama_model_params model_params = llama_model_default_params(); - model_params.devices = devices; - model_params.n_gpu_layers = 99; - // Keep AUTO so advertised device capabilities must enable mmap import. - if (model_params.load_mode != LLAMA_LOAD_MODE_AUTO) { - throw std::runtime_error("mmap import test requires the default AUTO load mode"); - } - llama_model_ptr imported(llama_model_load_from_file(path.string().c_str(), model_params)); - if (!imported) { - throw std::runtime_error("failed to load mmap-backed HRX model"); - } - // The loader retains the mapped bytes needed by imported buffers after the file is unlinked. - if (!std::filesystem::remove(path)) { - throw std::runtime_error("failed to unlink mmap-backed GGUF before readback"); - } - llama_context_params context_params = llama_context_default_params(); - context_params.n_ctx = 0; - context_params.n_ubatch = 64; - llama_context_ptr imported_context(llama_init_from_model(imported.get(), context_params)); - if (!imported_context) { - throw std::runtime_error("failed to initialize mmap-backed HRX context"); - } - const ggml_tensor * imported_weight = imported->get_tensor("blk.0.attn_q.weight"); - if (imported_weight == nullptr) { - throw std::runtime_error("mmap-backed HRX attention weight is missing"); - } - if (imported_weight->buffer == nullptr) { - throw std::runtime_error("mmap-backed HRX attention weight is unbound"); - } - ggml_backend_dev_t weight_device = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(imported_weight->buffer)); - if (weight_device != hrx_device) { - throw std::runtime_error("tiny attention weight was not placed on HRX"); - } - if (ggml_backend_buffer_is_host(imported_weight->buffer)) { - throw std::runtime_error("mmap-backed HRX attention weight was classified as host storage"); - } - if (!ggml_backend_hrx_buffer_is_imported(imported_weight->buffer)) { - throw std::runtime_error("tiny attention weight was not imported from mmap"); - } - std::vector actual(expected.size()); - ggml_backend_tensor_get(imported_weight, actual.data(), 0, actual.size()); - if (actual != expected) { - throw std::runtime_error("mmap-backed HRX weight differs after loader cleanup"); - } - const std::vector logits = get_logits(imported.get(), imported_context.get(), {1}); - if (logits.empty()) { - throw std::runtime_error("mmap-backed HRX model produced no logits"); - } - return 0; -#endif -} - int main(int argc, char ** argv) { // init the logger at max verbosity. filter with a custom callback respecting the user-configure verbosity common_log_set_verbosity_thold(LOG_LEVEL_DEBUG); @@ -1483,7 +1381,6 @@ int main(int argc, char ** argv) { float stdev = 0.1f; std::string out; const char * target_backend = nullptr; - bool hrx_import = false; int verbosity = LOG_LEVEL_ERROR; @@ -1551,8 +1448,6 @@ int main(int argc, char ** argv) { usage(argv); return 1; } - } else if (strcmp(argv[i], "--hrx-import") == 0) { - hrx_import = true; } else { LOG_ERR("%s: unknown argument: %s\n", __func__, argv[i]); usage(argv); @@ -1568,9 +1463,6 @@ int main(int argc, char ** argv) { init_normal_pool(seed, stdev); try { - if (hrx_import) { - return test_hrx_mmap_import(seed, stdev); - } if (!out.empty()) { return save_models(arch_filter, seed, stdev, verbosity, out); } From 733c6a5deb8de5156275a9b6e1911f7334114cc9 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 22:55:39 +0000 Subject: [PATCH 11/18] tests: drop obsolete HRX weight unloading coverage Assisted-by: Codex --- tests/hrx-ops-test.cpp | 575 +---------------------------------------- 1 file changed, 5 insertions(+), 570 deletions(-) diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 6ed99c6fbd0d..ebe425bf7e1a 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3225,374 +3225,6 @@ static void run_weight_persistence_sequence_case() { ggml_backend_free(cpu); } -static void run_weight_persistence_overlap_case() { - constexpr int64_t input_size = 256; - constexpr int64_t output_size = 64; - const char * old_flag = std::getenv("GGML_HRX_DISABLE_GRAPH_REPLAY"); - const bool had_flag = old_flag != nullptr; - const std::string old_flag_value = had_flag ? old_flag : ""; - REQUIRE(setenv("GGML_HRX_DISABLE_GRAPH_REPLAY", "0", 1) == 0); - - ggml_backend_t cpu = init_cpu_backend(); - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(cpu != nullptr && backend != nullptr); - auto * state = static_cast(backend->context); - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); - ggml_tensor * cpu_w0 = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * cpu_w1 = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * cpu_disjoint_w = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * hrx_w0 = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * hrx_w1 = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * hrx_disjoint_w = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); - ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); - ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); - ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_w0, cpu_ids); - ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_w0, hrx_ids); - ggml_tensor * cpu_disjoint_rows = ggml_get_rows(cpu_ctx, cpu_disjoint_w, cpu_ids); - ggml_tensor * hrx_disjoint_rows = ggml_get_rows(hrx_ctx, hrx_disjoint_w, hrx_ids); - ggml_tensor * cpu_mixed = ggml_mul_mat(cpu_ctx, cpu_w1, cpu_rows); - ggml_tensor * hrx_mixed = ggml_mul_mat(hrx_ctx, hrx_w1, hrx_rows); - ggml_tensor * cpu_packed = ggml_mul_mat(cpu_ctx, cpu_w1, cpu_input); - ggml_tensor * hrx_packed = ggml_mul_mat(hrx_ctx, hrx_w1, hrx_input); - ggml_cgraph * cpu_mixed_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * cpu_native_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * cpu_packed_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * cpu_disjoint_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_mixed_graph = ggml_new_graph(hrx_ctx); - ggml_cgraph * hrx_native_graph = ggml_new_graph(hrx_ctx); - ggml_cgraph * hrx_packed_graph = ggml_new_graph(hrx_ctx); - ggml_cgraph * hrx_disjoint_graph = ggml_new_graph(hrx_ctx); - ggml_build_forward_expand(cpu_mixed_graph, cpu_mixed); - ggml_build_forward_expand(cpu_native_graph, cpu_rows); - ggml_build_forward_expand(cpu_packed_graph, cpu_packed); - ggml_build_forward_expand(cpu_disjoint_graph, cpu_disjoint_rows); - ggml_build_forward_expand(hrx_mixed_graph, hrx_mixed); - ggml_build_forward_expand(hrx_native_graph, hrx_rows); - ggml_build_forward_expand(hrx_packed_graph, hrx_packed); - ggml_build_forward_expand(hrx_disjoint_graph, hrx_disjoint_rows); - hrx_mixed_graph->uid = 951071; - hrx_native_graph->uid = 951072; - hrx_packed_graph->uid = 951073; - hrx_disjoint_graph->uid = 951074; - - std::vector weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 31); - ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( - ggml_backend_get_device(backend), weight.data(), weight.size(), weight.size()); - REQUIRE(imported != nullptr); - REQUIRE(ggml_backend_tensor_alloc(imported, hrx_w0, weight.data()) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_tensor_alloc(imported, hrx_w1, weight.data()) == GGML_STATUS_SUCCESS); - std::vector disjoint_weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 43); - ggml_backend_buffer_t disjoint_imported = ggml_backend_dev_buffer_from_host_ptr( - ggml_backend_get_device(backend), disjoint_weight.data(), disjoint_weight.size(), disjoint_weight.size()); - REQUIRE(disjoint_imported != nullptr); - REQUIRE(ggml_backend_tensor_alloc(disjoint_imported, hrx_disjoint_w, disjoint_weight.data()) == GGML_STATUS_SUCCESS); - ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_buffer_set_usage(disjoint_imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); - ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); - REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); - ggml_backend_tensor_set(cpu_w0, weight.data(), 0, weight.size()); - ggml_backend_tensor_set(cpu_w1, weight.data(), 0, weight.size()); - ggml_backend_tensor_set(cpu_disjoint_w, disjoint_weight.data(), 0, disjoint_weight.size()); - const std::array ids = { 0, 1, 63 }; - set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); - const std::vector input = make_pattern_f32(input_size * 3, 29, 0.01f); - set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); - REQUIRE(ggml_backend_graph_compute(cpu, cpu_mixed_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(cpu, cpu_packed_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(cpu, cpu_disjoint_graph) == GGML_STATUS_SUCCESS); - const std::vector expected_mixed = get_f32_tensor(cpu, cpu_mixed); - const std::vector expected_native = get_f32_tensor(cpu, cpu_rows); - const std::vector expected_packed = get_f32_tensor(cpu, cpu_packed); - const std::vector expected_disjoint = get_f32_tensor(cpu, cpu_disjoint_rows); - REQUIRE(std::any_of(expected_mixed.begin(), expected_mixed.end(), - [](float value) { return std::fabs(value) > 0.1f; })); - REQUIRE(std::any_of(expected_packed.begin(), expected_packed.end(), - [](float value) { return std::fabs(value) > 0.1f; })); - - auto imported_graph = ggml::hrx::import_ggml_graph(*hrx_mixed_graph); - REQUIRE(imported_graph.valid()); - const auto * w0_value = imported_graph.graph.values().find_tensor(hrx_w0); - const auto * w1_value = imported_graph.graph.values().find_tensor(hrx_w1); - REQUIRE(w0_value != nullptr && w1_value != nullptr); - REQUIRE(w0_value->id != w1_value->id); - - // Materialize both layouts before gating real submissions; cold uploads wait on the stream. - REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - auto packed_lookup = state->graph_programs.get_or_build( - *hrx_packed_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - REQUIRE(packed_lookup.valid()); - auto packed_match = packed_lookup.program->match_trusted_graph(*hrx_packed_graph, true); - REQUIRE(packed_match.valid()); - std::vector packed_external; - for (const auto & external : packed_match.external_bindings) { - ggml::hrx::ValueBufferBinding buffer; - REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); - packed_external.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, ggml_nbytes(external.tensor) == 0)); - } - auto packed_bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(packed_external)); - REQUIRE(packed_bindings.valid()); - auto packed_demands = ggml::hrx::collect_command_program_weight_demands( - packed_lookup.program->commands(), packed_bindings); - REQUIRE(packed_demands.valid()); - REQUIRE(packed_demands.weights.size() == 1 && packed_demands.weights.front().requires_materialization); - auto preuploaded = state->host_weights.acquire(state->device->device, state->stream, - state->host_transfers, packed_demands.weights.front().source); - REQUIRE(preuploaded.valid()); - ggml_backend_hrx_cache_stats cross_before = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cross_before)); - REQUIRE(cross_before.weight_duplicate_warnings == 1); - // The public executors finish their kernel reads before returning. Queue warmed recordings directly - // to exercise ownership of genuinely outstanding reads without weakening that completion boundary. - struct Submission { - std::shared_ptr resources; - }; - auto prepare_submission = [&](ggml_cgraph * graph) { - auto lookup = state->graph_programs.get_or_build( - *graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - REQUIRE(lookup.valid()); - auto match = lookup.program->match_trusted_graph(*graph, true); - REQUIRE(match.valid()); - std::vector external_bindings; - for (const auto & external : match.external_bindings) { - ggml::hrx::ValueBufferBinding buffer; - REQUIRE(ggml_backend_hrx_resolve_value_buffer(external.tensor, buffer)); - external_bindings.push_back(ggml::hrx::make_command_program_binding( - external.value, buffer, ggml_nbytes(external.tensor) == 0)); - } - auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); - const ggml::hrx::CommandProgramExecutionContext context = { - state->device->device, state->stream, state->device->architecture.c_str(), - &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, - &state->host_transfers, &state->host_weights, &state->device->host_buffers, - &state->graph_replay_state, - }; - auto execution = state->prepared_programs.execute_with_result( - context, lookup.program->uid(), lookup.program->command_shape_hash(), lookup.program->commands(), bindings); - REQUIRE(execution.success); - REQUIRE(execution.execution_resources->recorded.valid()); - REQUIRE(execution.execution_resources->recorded.resident_host_weights == - execution.execution_resources->prepared.resident_host_weights); - Submission submission; - submission.resources = std::move(execution.execution_resources); - return submission; - }; - auto queue_submission = [&](Submission & submission) { - const auto launch = ggml::hrx::take_status(hrx_graph_exec_launch( - submission.resources->recorded.exec, state->stream)); - if (launch) { - return false; - } - submission.resources.reset(); - return !ggml::hrx::take_status(hrx_stream_flush(state->stream)); - }; - auto cross_native_submission = prepare_submission(hrx_native_graph); - auto cross_packed_submission = prepare_submission(hrx_packed_graph); - const std::vector poisoned_packed(expected_packed.size(), 123.0f); - ggml_backend_tensor_set(hrx_packed, poisoned_packed.data(), 0, poisoned_packed.size() * sizeof(float)); - hrx_semaphore_t cross_gate = nullptr; - REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &cross_gate))); - REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { cross_gate, 1 }))); - std::mutex cross_mutex; - std::condition_variable cross_cv; - bool cross_finished = false; - bool cross_timed_out = false; - std::thread cross_watchdog([&] { - std::unique_lock lock(cross_mutex); - if (!cross_cv.wait_for(lock, std::chrono::seconds(5), [&] { return cross_finished; })) { - cross_timed_out = true; - hrx_status_t signal = hrx_semaphore_signal(cross_gate, 1); - hrx_status_ignore(signal); - } - }); - const bool cross_native = queue_submission(cross_native_submission); - const bool cross_packed = cross_native && - queue_submission(cross_packed_submission); - ggml_backend_hrx_cache_stats cross_queued = {}; - const bool cross_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cross_queued); - { - std::lock_guard lock(cross_mutex); - cross_finished = true; - } - cross_cv.notify_one(); - const bool cross_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cross_gate, 1)); - cross_watchdog.join(); - REQUIRE(cross_native && cross_packed && cross_stats_ok && cross_release_ok && !cross_timed_out); - REQUIRE(cross_queued.weight_live_allocations == 2); - REQUIRE(cross_queued.weight_duplicate_warnings == cross_before.weight_duplicate_warnings); - preuploaded.lease = {}; - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); - hrx_semaphore_release(cross_gate); - - REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats warmed = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); - REQUIRE(warmed.weight_live_allocations == 2); - REQUIRE(warmed.weight_duplicate_warnings >= cross_queued.weight_duplicate_warnings); - // Graph execution and the direct submission helper share the warmed native entry. - REQUIRE(state->prepared_programs.stats().builds == 3); - REQUIRE(state->prepared_programs.stats().hits >= 1); - - // A cold preparation failure must leave an already queued native reader's lease intact. - REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - auto cold_native_submission = prepare_submission(hrx_native_graph); - const std::vector poisoned_native(expected_native.size(), 123.0f); - ggml_backend_tensor_set(hrx_rows, poisoned_native.data(), 0, - poisoned_native.size() * sizeof(float)); - hrx_semaphore_t cold_gate = nullptr; - REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &cold_gate))); - REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { cold_gate, 1 }))); - std::mutex cold_mutex; - std::condition_variable cold_cv; - bool cold_finished = false; - bool cold_timed_out = false; - std::thread cold_watchdog([&] { - std::unique_lock lock(cold_mutex); - if (!cold_cv.wait_for(lock, std::chrono::seconds(5), [&] { return cold_finished; })) { - cold_timed_out = true; - hrx_status_t signal = hrx_semaphore_signal(cold_gate, 1); - hrx_status_ignore(signal); - } - }); - const bool native_before_failure = - queue_submission(cold_native_submission); - state->host_weights.fail_cold_acquire_after_for_test(0); - const bool cold_failed = - ggml_backend_graph_compute_async(backend, hrx_disjoint_graph) == GGML_STATUS_FAILED; - ggml_backend_hrx_cache_stats cold_pending = {}; - const bool cold_stats_ok = ggml_backend_hrx_get_cache_stats(backend, &cold_pending); - { - std::lock_guard lock(cold_mutex); - cold_finished = true; - } - cold_cv.notify_one(); - const bool cold_release_ok = !ggml::hrx::take_status(hrx_semaphore_signal(cold_gate, 1)); - cold_watchdog.join(); - REQUIRE(native_before_failure && cold_failed && cold_stats_ok && cold_release_ok && !cold_timed_out); - REQUIRE(cold_pending.weight_live_allocations == 2); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); - hrx_semaphore_release(cold_gate); - REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &warmed)); - REQUIRE(warmed.weight_live_allocations == 2); - - auto mixed_submission = prepare_submission(hrx_mixed_graph); - auto native_submission = prepare_submission(hrx_native_graph); - auto packed_submission = prepare_submission(hrx_packed_graph); - const std::vector poisoned_mixed(expected_mixed.size(), 123.0f); - ggml_backend_tensor_set(hrx_mixed, poisoned_mixed.data(), 0, poisoned_mixed.size() * sizeof(float)); - ggml_backend_tensor_set(hrx_packed, poisoned_packed.data(), 0, poisoned_packed.size() * sizeof(float)); - hrx_semaphore_t gate = nullptr; - REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); - REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); - std::mutex gate_mutex; - std::condition_variable gate_cv; - bool finished = false; - bool timed_out = false; - std::thread watchdog([&] { - std::unique_lock lock(gate_mutex); - const bool released = gate_cv.wait_for(lock, std::chrono::seconds(5), [&] { return finished; }); - if (!released) { - timed_out = true; - hrx_status_t signal = hrx_semaphore_signal(gate, 1); - hrx_status_ignore(signal); - } - }); - const bool mixed_queued = queue_submission(mixed_submission); - const bool native_queued = mixed_queued && - queue_submission(native_submission); - const bool packed_queued = native_queued && - queue_submission(packed_submission); - bool stream_complete = true; - const bool stream_query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); - ggml_backend_hrx_cache_stats queued = {}; - const bool stats_ok = ggml_backend_hrx_get_cache_stats(backend, &queued); - { - std::lock_guard lock(gate_mutex); - finished = true; - } - gate_cv.notify_one(); - const bool signal_ok = !ggml::hrx::take_status(hrx_semaphore_signal(gate, 1)); - watchdog.join(); - REQUIRE(mixed_queued && native_queued && packed_queued); - REQUIRE(stream_query_ok && !stream_complete); - REQUIRE(stats_ok); - REQUIRE(queued.weight_live_allocations == 2); - REQUIRE(queued.weight_peak_bytes == warmed.weight_peak_bytes); - REQUIRE(queued.weight_duplicate_warnings == warmed.weight_duplicate_warnings); - REQUIRE(signal_ok); - REQUIRE(!timed_out); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); - require_close(get_f32_tensor(backend, hrx_packed), expected_packed, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats completed = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &completed)); - REQUIRE(completed.weight_live_allocations == 2); - REQUIRE(completed.weight_live_bytes == queued.weight_live_bytes); - REQUIRE(completed.weight_duplicate_warnings == queued.weight_duplicate_warnings); - REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_rows), expected_native, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats reused = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &reused)); - REQUIRE(reused.weight_materializations == completed.weight_materializations); - REQUIRE(reused.weight_live_allocations == 2); - - REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - REQUIRE(ggml_backend_graph_compute(backend, hrx_disjoint_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_disjoint_rows), expected_disjoint, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats mixed_and_disjoint = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &mixed_and_disjoint)); - REQUIRE(mixed_and_disjoint.weight_live_allocations == 3); - REQUIRE(ggml_backend_graph_compute(backend, hrx_mixed_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_mixed), expected_mixed, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats mixed_reused = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &mixed_reused)); - REQUIRE(mixed_reused.weight_materializations == mixed_and_disjoint.weight_materializations); - REQUIRE(mixed_reused.weight_live_allocations == 3); - REQUIRE(ggml_backend_graph_compute(backend, hrx_native_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - ggml_backend_hrx_cache_stats native_reused = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &native_reused)); - REQUIRE(native_reused.weight_live_allocations == 3); - REQUIRE(ggml_backend_graph_compute(backend, hrx_disjoint_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_disjoint_rows), expected_disjoint, 2.0e-2f, 3.0e-2f); - ggml_backend_hrx_cache_stats disjoint_reused = {}; - REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &disjoint_reused)); - REQUIRE(disjoint_reused.weight_materializations == native_reused.weight_materializations); - REQUIRE(disjoint_reused.weight_live_allocations == 3); - - hrx_semaphore_release(gate); - ggml_backend_buffer_free(hrx_tensors); - ggml_backend_buffer_free(cpu_tensors); - ggml_backend_buffer_free(imported); - ggml_backend_buffer_free(disjoint_imported); - ggml_free(hrx_ctx); - ggml_free(cpu_ctx); - ggml_backend_free(backend); - ggml_backend_free(cpu); - restore_environment_value("GGML_HRX_DISABLE_GRAPH_REPLAY", had_flag, old_flag_value); -} - static void run_failed_submission_weight_lifetime_case(bool uncached = false) { constexpr int64_t input_size = 256; constexpr int64_t output_size = 64; @@ -3613,21 +3245,15 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, input_size, output_size); ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, input_size, output_size); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); - ggml_tensor * hrx_native = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - ggml_cgraph * hrx_native_graph = ggml_new_graph(hrx_ctx); ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - ggml_build_forward_expand(hrx_native_graph, hrx_native); hrx_graph->uid = 951081; - hrx_native_graph->uid = 951082; std::vector weight = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 59); ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( @@ -3644,8 +3270,6 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); REQUIRE(cpu_tensors != nullptr && hrx_tensors != nullptr); ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); - const std::array ids = { 0, 1, 63 }; - ggml_backend_tensor_set(hrx_ids, ids.data(), 0, sizeof(ids)); const std::vector input = make_pattern_f32(input_size * 3, 41, 0.01f); set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); @@ -3672,22 +3296,14 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { }; auto packed_lookup = state->graph_programs.get_or_build( *hrx_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - auto native_lookup = state->graph_programs.get_or_build( - *hrx_native_graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - REQUIRE(packed_lookup.valid() && native_lookup.valid()); + REQUIRE(packed_lookup.valid()); auto packed_bindings = bind(*packed_lookup.program, *hrx_graph); - auto native_bindings = bind(*native_lookup.program, *hrx_native_graph); auto packed_demands = ggml::hrx::collect_command_program_weight_demands( packed_lookup.program->commands(), packed_bindings); - auto native_demands = ggml::hrx::collect_command_program_weight_demands( - native_lookup.program->commands(), native_bindings); - REQUIRE(packed_demands.valid() && native_demands.valid()); + REQUIRE(packed_demands.valid()); REQUIRE(packed_demands.weights.size() == 1 && packed_demands.weights.front().requires_materialization); - REQUIRE(native_demands.weights.size() == 1); const auto packed_source = packed_demands.weights.front().source; - const auto native_source = native_demands.weights.front().source; - REQUIRE(native_source.layout == ggml::hrx::kNativeWeightLayout); - REQUIRE(packed_source.layout != native_source.layout); + REQUIRE(packed_source.layout != ggml::hrx::kNativeWeightLayout); REQUIRE(packed_source.materialized_length > 0); auto preuploaded = state->host_weights.acquire(state->device->device, state->stream, state->host_transfers, packed_source); @@ -4110,49 +3726,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, } } REQUIRE(selected_row64); - if (coherent_subregion) { - // This internal factory must preserve the owner CPU address even when the import binds directly by handle. - auto graph_import = ggml::hrx::import_ggml_graph(*hrx_decode); - REQUIRE(graph_import.valid()); - auto & values = graph_import.graph.values(); - for (ggml::hrx::ValueId id : values.external_value_ids()) { - const ggml::hrx::Value * value = values.find(id); - REQUIRE(value != nullptr); - REQUIRE(value->tensor != nullptr); - ggml::hrx::ValueBufferBinding source_binding; - REQUIRE(ggml_backend_hrx_resolve_value_buffer(value->tensor, source_binding)); - REQUIRE(values.bind_buffer(id, source_binding)); - } - const auto * weight_value = values.find_tensor(hrx_weight); - REQUIRE(weight_value != nullptr); - auto value_map_bindings = ggml::hrx::CommandProgramBindings::from_value_map(values); - REQUIRE(value_map_bindings.valid()); - const auto * weight_binding = value_map_bindings.find(weight_value->id); - REQUIRE(weight_binding != nullptr); - REQUIRE(weight_binding->buffer != nullptr); - REQUIRE(weight_binding->host_data == nullptr); - REQUIRE(weight_binding->source_host_data == coherent_owner_base); - REQUIRE(weight_binding->source_offset == 256); - REQUIRE(weight_binding->source_capacity == source.size() + 256); - const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); - const ggml::hrx::CommandProgramExecutionContext execution_context = { - backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), - &corpus, &backend_state->kernel_executables, &backend_state->transient_arena, - &backend_state->host_transfers, &backend_state->host_weights, &backend_state->device->host_buffers, - &backend_state->graph_replay_state, - }; - auto prepared = ggml::hrx::prepare_command_program( - execution_context, lookup.program->commands(), value_map_bindings); - REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights->size() == 1); - ggml::hrx::RecordedCommandGraph recorded; - const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( - execution_context, lookup.program->commands(), value_map_bindings, prepared, recorded); - REQUIRE(launch.success); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_decode_output), - get_f32_tensor(cpu_backend, cpu_decode_output), 1.0f, 3.0e-2f); - } compute_and_compare(cpu_other, hrx_other, cpu_other_output, hrx_other_output); REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); REQUIRE(stats.weight_materializations == 1); @@ -4199,7 +3772,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, compute_and_compare(cpu_native, hrx_native, cpu_native_output, hrx_native_output); REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); - const uint64_t native_materializations = coherent_subregion ? 1 : 2; + const uint64_t native_materializations = 2; REQUIRE(stats.weight_materializations == native_materializations); REQUIRE(stats.weight_live_allocations == native_materializations); const uint64_t uploaded_bytes = stats.weight_uploaded_bytes; @@ -4272,7 +3845,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, REQUIRE(after_rebind.prepared_program_builds == before_rebind.prepared_program_builds + 1); REQUIRE(after_rebind.weight_materializations == before_rebind.weight_materializations); - if (source_kind == PackedWeightSourceKind::Imported) { + { // Exercise production preparation and recording: physical residency must remain counted after both // the graph-program owner and weight lookup entries are gone, for as long as recording retains it. auto full_match = lookup.program->match_trusted_graph(*hrx_decode, true); @@ -4321,9 +3894,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, ggml_backend_buffer_free(decode_host); ggml_backend_buffer_free(rebound_host); ggml_backend_buffer_free(imported); - if (owner != nullptr) { - ggml_backend_buffer_free(owner); - } ggml_free(cpu_ctx); ggml_free(hrx_ctx); ggml_backend_free(cpu_backend); @@ -4629,134 +4199,6 @@ static void run_mixed_weight_layout_program_case(ggml_type weight_type) { ggml_backend_free(cpu); } -static void run_native_cpu_weight_during_import_layout_transition_case(ggml_type weight_type) { - constexpr int64_t input_size = 256; - constexpr int64_t output_size = 64; - ggml_backend_t cpu = init_cpu_backend(); - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(cpu != nullptr && backend != nullptr); - auto * state = static_cast(backend->context); - - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); - - ggml_tensor * cpu_host_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); - ggml_tensor * cpu_imported_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, 3); - ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, 3); - ggml_tensor * hrx_host_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); - ggml_tensor * hrx_imported_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, 3); - ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, 3); - ggml_tensor * cpu_host_rows = ggml_get_rows(cpu_ctx, cpu_host_weight, cpu_ids); - ggml_tensor * hrx_host_rows = ggml_get_rows(hrx_ctx, hrx_host_weight, hrx_ids); - ggml_tensor * cpu_imported_rows = ggml_get_rows(cpu_ctx, cpu_imported_weight, cpu_ids); - ggml_tensor * hrx_imported_rows = ggml_get_rows(hrx_ctx, hrx_imported_weight, hrx_ids); - ggml_tensor * cpu_imported_matmul = ggml_mul_mat(cpu_ctx, cpu_imported_weight, cpu_input); - ggml_tensor * hrx_imported_matmul = ggml_mul_mat(hrx_ctx, hrx_imported_weight, hrx_input); - ggml_cgraph * cpu_packed = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_packed = ggml_new_graph(hrx_ctx); - ggml_cgraph * cpu_native = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_native = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_packed != nullptr && hrx_packed != nullptr && cpu_native != nullptr && hrx_native != nullptr); - ggml_build_forward_expand(cpu_packed, cpu_host_rows); - ggml_build_forward_expand(cpu_packed, cpu_imported_matmul); - ggml_build_forward_expand(hrx_packed, hrx_host_rows); - ggml_build_forward_expand(hrx_packed, hrx_imported_matmul); - ggml_build_forward_expand(cpu_native, cpu_host_rows); - ggml_build_forward_expand(cpu_native, cpu_imported_rows); - ggml_build_forward_expand(hrx_native, hrx_host_rows); - ggml_build_forward_expand(hrx_native, hrx_imported_rows); - hrx_packed->uid = 951061; - hrx_native->uid = 951062; - - std::vector host_bytes = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); - std::vector imported_bytes = make_matmul_weight_bytes(weight_type, input_size, output_size, 11); - ggml_backend_buffer_t host_buffer = ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), host_bytes.size()); - REQUIRE(host_buffer != nullptr); - REQUIRE(ggml_backend_tensor_alloc(host_buffer, hrx_host_weight, ggml_backend_buffer_get_base(host_buffer)) == - GGML_STATUS_SUCCESS); - ggml_backend_buffer_set_usage(host_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_tensor_set(hrx_host_weight, host_bytes.data(), 0, host_bytes.size()); - ggml_backend_buffer_t imported = ggml_backend_dev_buffer_from_host_ptr( - ggml_backend_get_device(backend), imported_bytes.data(), imported_bytes.size(), imported_bytes.size()); - REQUIRE(imported != nullptr); - REQUIRE(ggml_backend_tensor_alloc(imported, hrx_imported_weight, imported_bytes.data()) == GGML_STATUS_SUCCESS); - ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_buffer_t hrx_tensors = ggml_backend_alloc_ctx_tensors(hrx_ctx, backend); - ggml_backend_buffer_t cpu_tensors = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); - REQUIRE(hrx_tensors != nullptr && cpu_tensors != nullptr); - ggml_backend_tensor_set(cpu_host_weight, host_bytes.data(), 0, host_bytes.size()); - ggml_backend_tensor_set(cpu_imported_weight, imported_bytes.data(), 0, imported_bytes.size()); - const std::array ids = { 1, 17, 63 }; - const std::vector input = make_pattern_f32(input_size * 3, 7, 0.01f); - set_tensor_pair_bytes(cpu, cpu_ids, backend, hrx_ids, ids.data(), sizeof(ids)); - set_tensor_pair_bytes(cpu, cpu_input, backend, hrx_input, input.data(), input.size() * sizeof(float)); - - const std::string packed_layout = weight_type == GGML_TYPE_Q4_K ? - ggml::hrx::kQ4KPackedK256Row64Layout : ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; - auto check_layouts = [&](ggml_cgraph * graph, bool expect_packed) { - auto lookup = state->graph_programs.get_or_build( - *graph, ggml::hrx::get_qwen_kernel_corpus(), state->device->architecture); - REQUIRE(lookup.valid()); - const auto * host_value = lookup.program->graph().values().find_tensor(hrx_host_weight); - const auto * imported_value = lookup.program->graph().values().find_tensor(hrx_imported_weight); - REQUIRE(host_value != nullptr && imported_value != nullptr); - bool host_native = false; - bool imported_selected = false; - for (const auto & command : lookup.program->commands().commands) { - for (const auto & binding : command.bindings) { - if (binding.value == host_value->id) { - REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); - host_native = true; - } - if (binding.value == imported_value->id) { - const std::string & expected = expect_packed ? packed_layout : ggml::hrx::kNativeWeightLayout; - imported_selected = imported_selected || binding.layout == expected; - } - } - } - REQUIRE(host_native && imported_selected); - }; - check_layouts(hrx_packed, true); - check_layouts(hrx_native, false); - - auto compute_and_compare = [&](ggml_cgraph * cpu_graph, ggml_cgraph * hrx_graph, - ggml_tensor * cpu_imported_output, ggml_tensor * hrx_imported_output) { - REQUIRE(ggml_backend_graph_compute(cpu, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - require_close(get_f32_tensor(backend, hrx_host_rows), get_f32_tensor(cpu, cpu_host_rows), 2.0e-3f); - require_close(get_f32_tensor(backend, hrx_imported_output), - get_f32_tensor(cpu, cpu_imported_output), 1.0f, 3.0e-2f); - }; - compute_and_compare(cpu_packed, hrx_packed, cpu_imported_matmul, hrx_imported_matmul); - auto first = state->host_weights.stats(); - REQUIRE(first.misses == 2); - REQUIRE(first.allocation_count == 2); - compute_and_compare(cpu_native, hrx_native, cpu_imported_rows, hrx_imported_rows); - auto native = state->host_weights.stats(); - REQUIRE(native.misses == first.misses + 1); - REQUIRE(native.allocation_count == 3); - compute_and_compare(cpu_packed, hrx_packed, cpu_imported_matmul, hrx_imported_matmul); - auto packed_again = state->host_weights.stats(); - REQUIRE(packed_again.misses == native.misses); - REQUIRE(packed_again.allocation_count == 3); - - ggml_backend_buffer_free(hrx_tensors); - ggml_backend_buffer_free(cpu_tensors); - ggml_backend_buffer_free(imported); - ggml_backend_buffer_free(host_buffer); - ggml_free(hrx_ctx); - ggml_free(cpu_ctx); - ggml_backend_free(backend); - ggml_backend_free(cpu); -} - static void run_q4_prefill_source_eligibility_case() { constexpr int64_t input_size = 256; constexpr int64_t output_size = 64; @@ -9088,7 +8530,6 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("mutable_host_input_after_writeback.cpu", [] { run_mutable_host_input_after_writeback_case(false); }); suite.device_case("mutable_host_input_after_writeback.hrx", [] { run_mutable_host_input_after_writeback_case(true); }); suite.device_case("weight_persistence_sequence", [] { run_weight_persistence_sequence_case(); }); - suite.device_case("weight_persistence_overlap", [] { run_weight_persistence_overlap_case(); }); suite.device_case("failed_submission_weight_lifetime", [] { run_failed_submission_weight_lifetime_case(); }); suite.device_case("failed_submission_weight_lifetime.uncached", [] { run_failed_submission_weight_lifetime_case(true); }); suite.device_case("imported_packed_weight.q4_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q4_K); }); @@ -9108,12 +8549,6 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { suite.device_case("coherent_subregion_native_weight.q4_k", [] { run_imported_packed_weight_case(GGML_TYPE_Q4_K, PackedWeightSourceKind::CoherentSubregion); }); - suite.device_case("native_cpu_weight_import_transition.q4_k", [] { - run_native_cpu_weight_during_import_layout_transition_case(GGML_TYPE_Q4_K); - }); - suite.device_case("native_cpu_weight_import_transition.q6_k", [] { - run_native_cpu_weight_during_import_layout_transition_case(GGML_TYPE_Q6_K); - }); suite.device_case("q4_prefill_source_eligibility", [] { run_q4_prefill_source_eligibility_case(); }); suite.device_case("uploaded_native_q6_vector", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128, From 405a1944f8342581952b67969da4a444cf9282cb Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 23:15:18 +0000 Subject: [PATCH 12/18] ggml-hrx: remove obsolete weight retirement machinery Assisted-by: Codex --- ggml/include/ggml-hrx.h | 2 +- ggml/src/ggml-hrx/README.md | 8 +- ggml/src/ggml-hrx/backend-buffer-binding.cpp | 50 +---- ggml/src/ggml-hrx/backend-buffer-binding.h | 9 - ggml/src/ggml-hrx/ggml-hrx.cpp | 42 ---- .../runtime/command-program-executor.cpp | 209 ++++++++---------- .../runtime/command-program-executor.h | 27 +-- ggml/src/ggml-hrx/runtime/graph-executor.cpp | 4 - .../ggml-hrx/runtime/graph-program-cache.cpp | 43 +--- .../ggml-hrx/runtime/graph-program-cache.h | 3 - ggml/src/ggml-hrx/runtime/host-memory.cpp | 144 ++---------- ggml/src/ggml-hrx/runtime/host-memory.h | 26 +-- .../prepared-command-program-cache.cpp | 104 +++------ .../runtime/prepared-command-program-cache.h | 6 +- tests/hrx-backend-test.cpp | 71 +++--- tests/hrx-buffer-test.cpp | 22 +- tests/hrx-ops-test.cpp | 51 +++-- 17 files changed, 226 insertions(+), 595 deletions(-) diff --git a/ggml/include/ggml-hrx.h b/ggml/include/ggml-hrx.h index 999160950b3b..2fb96acc36e4 100644 --- a/ggml/include/ggml-hrx.h +++ b/ggml/include/ggml-hrx.h @@ -14,6 +14,7 @@ struct ggml_backend_hrx_cache_stats { uint64_t weight_materializations; uint64_t weight_reuses; uint64_t weight_uploaded_bytes; + // Current weight-cache residency; peak bytes and materializations are cumulative. uint64_t weight_live_allocations; uint64_t weight_live_bytes; uint64_t weight_peak_bytes; @@ -26,7 +27,6 @@ GGML_BACKEND_API ggml_backend_t ggml_backend_hrx_init(size_t device) GGML_BACKEND_API bool ggml_backend_is_hrx(ggml_backend_t backend); GGML_BACKEND_API int ggml_backend_hrx_get_device_count(void); GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device); -GGML_BACKEND_API bool ggml_backend_hrx_buffer_is_imported(ggml_backend_buffer_t buffer); GGML_BACKEND_API bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, struct ggml_backend_hrx_cache_stats * stats); GGML_BACKEND_API ggml_backend_reg_t ggml_backend_hrx_reg(void); diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index ef71302730b8..988f94c3b836 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -4,13 +4,13 @@ Weight residency supports standard GGUF inference with immutable model weights. Alternative layouts require an HRX `buffer_from_host_ptr` import that does not alias a known HRX host allocation. Other weights remain native-only: host weights reuse a cached native copy, and directly bindable device weights reuse their existing allocation. Source availability for native caching does not grant permission to select a transformed layout. -Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx: weight_stats` line when each HRX backend is released. The `-v` flag enables the backend log callback in `llama-bench`. The line reports successful weight materializations, cache acquisition reuse, uploaded weight bytes, live allocation count and bytes, peak live bytes, and all host transfer uploads. `ggml_backend_hrx_get_cache_stats` exposes the same weight counters while a backend is running. Prepared program hits do not count as weight cache acquisitions, so unchanged replay can reuse a weight without increasing the `reused` counter. +Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx: weight_stats` line when each HRX backend is released. The `-v` flag enables the backend log callback in `llama-bench`. The line reports successful weight materializations, cache acquisition reuse, uploaded weight bytes, cached allocation count and bytes, peak cached bytes, and all host transfer uploads. `ggml_backend_hrx_get_cache_stats` exposes the same weight counters while a backend is running. Prepared program hits do not count as weight cache acquisitions, so unchanged replay can reuse a weight without increasing the `reused` counter. -Live weight bytes count physical allocations retained by the cache, a prepared program, or a recorded graph. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. +Resident weight bytes count physical allocations owned by the weight cache. An explicit clear resets the resident count and bytes while preserving cumulative materializations, reuse, warnings and peak bytes. Existing prepared leases and recorded buffer handles remain valid after a clear, but are no longer counted as cache residency. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. -Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared and recorded programs share allocation leases with the cache. +Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared programs hold allocation leases; recorded graphs retain the HRX buffer handles they use. -A graph instance remembers its validated prepared-cache entry. Its first execution resolves all physical bindings; repeated executions with the same graph UID and node array bind only mutable host staging. Structurally equivalent graphs can share command programs without sharing unchecked bindings: each new instance goes through the full binding lookup before using the fast path. Mutable source rebinding also returns to that lookup. Both execution paths use the same prepared-cache entry, lock, and replay implementation; warm executions do not rebuild per-weight demand. +A graph instance remembers its validated prepared-cache entry. Its first execution resolves all physical bindings; repeated executions with the same graph UID and node array bind only mutable host staging. Structurally equivalent graphs can share command programs without sharing unchecked bindings: each new instance goes through the full binding lookup before using the fast path. Mutable source rebinding also returns to that lookup. Both execution paths use the same prepared-cache entry, lock, and replay implementation; warm executions do not rebuild per-weight demand. On a preparation miss, one command-access analysis covers staging and weight demand. Every layout request is validated against the current source before equivalent demands are deduplicated. Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning groups exact source byte ranges; it does not detect partial overlap or infer a whole tensor from a transformed subrange. It names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp index 5e21c06e1dc8..10d03c2e04d5 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.cpp +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -3,39 +3,6 @@ #include "ggml-backend-impl.h" #include "ggml.h" -#include -#include - -namespace { - -std::mutex foreign_host_generations_mutex; -// GGML does not notify HRX when a foreign backend destroys a buffer. The source lifetime contract holds while an -// HRX backend is alive, and the table is retired with the last such backend. -std::unordered_map foreign_host_generations; -size_t active_hrx_backends = 0; - -uint64_t foreign_host_generation(uint64_t identity) { - std::lock_guard lock(foreign_host_generations_mutex); - const auto found = foreign_host_generations.find(identity); - return found == foreign_host_generations.end() ? 1 : found->second; -} - -} // namespace - -void ggml_backend_hrx_foreign_host_backend_started() { - std::lock_guard lock(foreign_host_generations_mutex); - ++active_hrx_backends; -} - -void ggml_backend_hrx_foreign_host_backend_stopped() { - std::lock_guard lock(foreign_host_generations_mutex); - GGML_ASSERT(active_hrx_backends > 0); - --active_hrx_backends; - if (active_hrx_backends == 0) { - foreign_host_generations.clear(); - } -} - ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer) { return static_cast(buffer->context); } @@ -104,7 +71,7 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx binding.length = ggml_nbytes(tensor); binding.identity = buffer_address ^ (base_address + 0x9e3779b97f4a7c15ull + (buffer_address << 6) + (buffer_address >> 2)); - binding.generation = foreign_host_generation(binding.identity); + binding.generation = 1; binding.capacity = capacity; binding.weight = ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; // Standard GGUF weights are immutable, so their native copy can use the shared weight cache. @@ -164,24 +131,9 @@ bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backe region.offset = binding.offset; } region.length = ggml_nbytes(tensor); - region.type = tensor->type; - for (size_t i = 0; i < GGML_MAX_DIMS; ++i) { - region.ne[i] = tensor->ne[i]; - region.nb[i] = tensor->nb[i]; - } return true; } -void ggml_backend_hrx_mark_foreign_host_region_written(const ggml_tensor * tensor) { - ggml_backend_hrx_imported_region region = {}; - if (!ggml_backend_hrx_get_retained_region(tensor, region)) { - return; - } - std::lock_guard lock(foreign_host_generations_mutex); - uint64_t & generation = foreign_host_generations[region.identity]; - generation = generation == 0 ? 2 : generation + 1; -} - bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region) { ggml_backend_hrx_buffer_context * context = nullptr; size_t offset = 0; diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.h b/ggml/src/ggml-hrx/backend-buffer-binding.h index bd110f4abe13..bd2935a36e66 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.h +++ b/ggml/src/ggml-hrx/backend-buffer-binding.h @@ -4,7 +4,6 @@ #include "ggml-backend.h" #include "graph/value-map.h" -#include #include struct ggml_tensor; @@ -16,9 +15,6 @@ struct ggml_backend_hrx_imported_region { uint64_t generation; size_t offset; size_t length; - ggml_type type; - std::array ne; - std::array nb; }; ggml_backend_hrx_buffer_context * ggml_backend_hrx_buffer_context_from_buffer(ggml_backend_buffer_t buffer); @@ -34,8 +30,3 @@ bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backe bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor); bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor); bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region); -void ggml_backend_hrx_mark_imported_region_written(const ggml_tensor * tensor); -void ggml_backend_hrx_mark_retained_region_written(const ggml_tensor * tensor); -void ggml_backend_hrx_mark_foreign_host_region_written(const ggml_tensor * tensor); -void ggml_backend_hrx_foreign_host_backend_started(); -void ggml_backend_hrx_foreign_host_backend_stopped(); diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 563354bc36f6..481af2c8f678 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -422,7 +422,6 @@ static void backend_free(ggml_backend_t backend) { hrx_stream_release(context->stream); delete context; delete backend; - ggml_backend_hrx_foreign_host_backend_stopped(); } static bool synchronous_upload_fallback(ggml_backend_hrx_context * backend, @@ -668,7 +667,6 @@ static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parame delete context; hrx_stream_release(stream); } else { - ggml_backend_hrx_foreign_host_backend_started(); } return backend; } @@ -1375,44 +1373,4 @@ ggml_backend_buffer_type_t ggml_backend_hrx_buffer_type(size_t device) { nullptr; } -bool ggml_backend_hrx_buffer_is_imported(ggml_backend_buffer_t buffer) { - if (buffer == nullptr) { - return false; - } - const auto * registry = static_cast(ggml_backend_hrx_reg()->context); - for (const auto & device : registry->device_contexts) { - const bool regular_type = buffer->buft == &device->buft; - const bool host_type = buffer->buft == &device->host_buft; - const bool hrx_type = regular_type || host_type; - if (hrx_type) { - return buffer_context(buffer)->imported; - } - } - return false; -} - -void ggml_backend_hrx_mark_imported_region_written(const ggml_tensor * tensor) { - ggml_backend_hrx_buffer_context * context = nullptr; - size_t offset = 0; - if (ggml_backend_hrx_tensor_binding(tensor, &context, &offset)) { - if (context->imported) { - ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); - } - } -} - -void ggml_backend_hrx_mark_retained_region_written(const ggml_tensor * tensor) { - ggml_backend_hrx_buffer_context * context = nullptr; - size_t offset = 0; - if (!ggml_backend_hrx_tensor_binding(tensor, &context, &offset)) { - ggml_backend_hrx_mark_foreign_host_region_written(tensor); - return; - } - ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; - const bool retained = context->imported || ggml_backend_buffer_is_host(buffer); - if (retained) { - ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); - } -} - GGML_BACKEND_DL_IMPL(ggml_backend_hrx_reg) diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index d6160f9e00cf..9348ad3a9bc0 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -111,7 +111,6 @@ RecordedCommandGraph::RecordedCommandGraph(RecordedCommandGraph && other) noexce graph(std::exchange(other.graph, nullptr)), exec(std::exchange(other.exec, nullptr)), retained_buffers(std::move(other.retained_buffers)), - resident_host_weights(std::move(other.resident_host_weights)), retained_executables(std::move(other.retained_executables)), bound_transient_arena_allocation_id(other.bound_transient_arena_allocation_id), dispatch_count(other.dispatch_count), @@ -137,7 +136,6 @@ RecordedCommandGraph & RecordedCommandGraph::operator=(RecordedCommandGraph && o graph = std::exchange(other.graph, nullptr); exec = std::exchange(other.exec, nullptr); retained_buffers = std::move(other.retained_buffers); - resident_host_weights = std::move(other.resident_host_weights); retained_executables = std::move(other.retained_executables); bound_transient_arena_allocation_id = other.bound_transient_arena_allocation_id; dispatch_count = other.dispatch_count; @@ -750,52 +748,109 @@ static bool same_weight_request(const CommandBinding & lhs, const CommandBinding return same_value && same_layout && same_transformation; } -static void append_weight_requests(const std::vector & commands, - const std::unordered_map & access_by_value, - std::unordered_map> & groups_by_value, - CommandProgramWeightRequests & result) { - for (const Command & command : commands) { - for (const CommandBinding & request : command.bindings) { - const bool graph_value = request.origin == CommandBindingOrigin::GraphValue; - const bool reads = request.access != ResourceAccess::Write; - const bool eligible_access = graph_value && reads; - if (!eligible_access) { - continue; - } - const auto found_access = access_by_value.find(request.value.value); - const bool value_has_access = found_access != access_by_value.end(); - const bool value_is_written = value_has_access && found_access->second.write; - if (value_is_written) { - continue; - } - auto & groups = groups_by_value[request.value.value]; - const auto found_group = std::find_if(groups.begin(), groups.end(), [&](size_t index) { - return same_weight_request(result.requests[index].binding, request); - }); - size_t group; - if (found_group == groups.end()) { - group = result.group_count++; - groups.push_back(result.requests.size()); - } else { - group = result.requests[*found_group].group; +static CommandProgramWeightDemands collect_weight_demands( + const CommandProgram & commands, + const CommandProgramBindings & bindings, + const std::unordered_map & access_by_value) { + CommandProgramWeightDemands result; + std::unordered_map bindings_by_value; + bindings_by_value.reserve(bindings.bindings().size()); + for (const CommandProgramBinding & binding : bindings.bindings()) { + bindings_by_value.try_emplace(binding.value.value, &binding); + } + std::unordered_map> demands_by_value; + auto collect = [&](const std::vector & command_list) { + for (const Command & command : command_list) { + for (const CommandBinding & request : command.bindings) { + const bool graph_value = request.origin == CommandBindingOrigin::GraphValue; + const bool reads = request.access != ResourceAccess::Write; + const bool eligible_access = graph_value && reads; + if (!eligible_access) { + continue; + } + const auto found_access = access_by_value.find(request.value.value); + const bool value_has_access = found_access != access_by_value.end(); + const bool value_is_written = value_has_access && found_access->second.write; + if (value_is_written) { + continue; + } + const auto found_binding = bindings_by_value.find(request.value.value); + const CommandProgramBinding * binding = + found_binding != bindings_by_value.end() ? found_binding->second : nullptr; + const bool bound_weight = binding != nullptr && binding->weight; + if (!bound_weight) { + continue; + } + const bool transformed = request.layout != kNativeWeightLayout; + const bool retained_source = binding->source_host_data != nullptr; + const bool missing_transformed_source = transformed && !retained_source; + if (missing_transformed_source) { + result.status.log("weight value %d has no retained source for transformed layout %s", + request.value.value, request.layout.c_str()); + continue; + } + if (!retained_source) { + continue; + } + if (transformed) { + const bool has_type = request.source_type != GGML_TYPE_COUNT; + const bool has_input_size = request.input_size > 0; + const bool has_output_size = request.output_size > 0; + const bool valid_shape = has_type && has_input_size && has_output_size; + const bool starts_at_source = request.offset == 0; + const bool has_source_length = request.source_length != 0; + const bool source_fits = request.source_length <= binding->length; + const bool has_output_length = request.length != 0; + const bool valid_range = starts_at_source && has_source_length && source_fits && has_output_length; + const bool valid_transformation = valid_shape && valid_range; + if (!valid_transformation) { + result.status.log("weight value %d has invalid transformed layout request %s", + request.value.value, request.layout.c_str()); + continue; + } + } + + // Validate every occurrence before deduplication, including malformed duplicates. + auto & indices = demands_by_value[request.value.value]; + const auto found = std::find_if(indices.begin(), indices.end(), [&](size_t index) { + return same_weight_request(result.weights[index].request, request); + }); + if (found != indices.end()) { + continue; + } + CommandProgramWeightDemand demand; + demand.request = request; + demand.requires_materialization = binding->requires_materialization() || transformed; + demand.source.host_data = binding->source_host_data; + demand.source.identity = binding->identity; + demand.source.generation = binding->generation; + demand.source.capacity = binding->source_capacity; + demand.source.offset = binding->source_offset; + demand.source.length = transformed ? request.source_length : binding->length; + demand.source.materialized_length = transformed ? request.length : binding->length; + demand.source.layout = request.layout; + demand.source.source_type = transformed ? request.source_type : GGML_TYPE_COUNT; + demand.source.input_size = request.input_size; + demand.source.output_size = request.output_size; + indices.push_back(result.weights.size()); + result.weights.push_back(std::move(demand)); } - result.requests.push_back({ request, group }); } - } + }; + collect(commands.initialization_commands); + collect(commands.commands); + return result; } static CommandProgramBindings materialize_host_bindings(const CommandProgramExecutionContext & context, const CommandProgram & commands, const CommandProgramBindings & bindings, PreparedCommandProgram & prepared) { - // Publish the owner before acquisition so partial preparation retains every acquired lease. - auto resident_weights = std::make_shared>(); - prepared.resident_host_weights = resident_weights; std::vector materialized; Status status; materialized.reserve(bindings.bindings().size()); const std::unordered_map access_by_value = collect_graph_value_access(commands); - const CommandProgramWeightDemands demands = collect_command_program_weight_demands(commands, bindings); + const CommandProgramWeightDemands demands = collect_weight_demands(commands, bindings, access_by_value); status.append(demands.status); for (const CommandProgramBinding & binding : bindings.bindings()) { const auto found_access = access_by_value.find(binding.value.value); @@ -837,7 +892,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec device_binding.offset = 0; device_binding.length = resident.lease.length(); device_binding.capacity = resident.lease.length(); - resident_weights->push_back(std::move(resident.lease)); + prepared.resident_host_weights.push_back(std::move(resident.lease)); } materialized.push_back(std::move(device_binding)); } @@ -1655,8 +1710,6 @@ static void retain_prepared_command_list_resources(const std::vector retained_buffers; std::unordered_map retained_executables; if (allocation.buffer != nullptr) { @@ -1788,85 +1841,9 @@ static bool execute_prepared_command_prefix_via_graph(const CommandProgramExecut } // namespace -CommandProgramWeightRequests build_command_program_weight_requests(const CommandProgram & commands) { - CommandProgramWeightRequests result; - std::unordered_map> groups_by_value; - const auto access_by_value = collect_graph_value_access(commands); - append_weight_requests(commands.initialization_commands, access_by_value, groups_by_value, result); - append_weight_requests(commands.commands, access_by_value, groups_by_value, result); - return result; -} - -CommandProgramWeightDemands bind_command_program_weight_demands(const CommandProgramWeightRequests & requests, - const CommandProgramBindings & bindings) { - CommandProgramWeightDemands result; - std::unordered_map bindings_by_value; - bindings_by_value.reserve(bindings.bindings().size()); - for (const CommandProgramBinding & binding : bindings.bindings()) { - bindings_by_value.try_emplace(binding.value.value, &binding); - } - std::vector emitted_groups(requests.group_count, 0); - result.weights.reserve(requests.group_count); - for (const CommandProgramWeightRequest & entry : requests.requests) { - const CommandBinding & request = entry.binding; - const auto found_binding = bindings_by_value.find(request.value.value); - const CommandProgramBinding * binding = - found_binding != bindings_by_value.end() ? found_binding->second : nullptr; - const bool bound_weight = binding != nullptr && binding->weight; - if (!bound_weight) { - continue; - } - const bool transformed = request.layout != kNativeWeightLayout; - const bool retained_source = binding->source_host_data != nullptr; - const bool missing_transformed_source = transformed && !retained_source; - if (missing_transformed_source) { - result.status.log("weight value %d has no retained source for transformed layout %s", - request.value.value, request.layout.c_str()); - continue; - } - if (!retained_source) { - continue; - } - if (transformed) { - const bool valid_shape = request.source_type != GGML_TYPE_COUNT && request.input_size > 0 && - request.output_size > 0; - const bool valid_range = request.offset == 0 && request.source_length != 0 && - request.source_length <= binding->length && request.length != 0; - const bool valid_transformation = valid_shape && valid_range; - if (!valid_transformation) { - result.status.log("weight value %d has invalid transformed layout request %s", - request.value.value, request.layout.c_str()); - continue; - } - } - // Validation precedes deduplication: an invalid duplicate must still report its error, - // and an invalid first occurrence must not suppress a later valid request. - if (emitted_groups[entry.group]) { - continue; - } - emitted_groups[entry.group] = 1; - CommandProgramWeightDemand demand; - demand.request = request; - demand.requires_materialization = binding->requires_materialization() || transformed; - demand.source.host_data = binding->source_host_data; - demand.source.identity = binding->identity; - demand.source.generation = binding->generation; - demand.source.capacity = binding->source_capacity; - demand.source.offset = binding->source_offset; - demand.source.length = transformed ? request.source_length : binding->length; - demand.source.materialized_length = transformed ? request.length : binding->length; - demand.source.layout = request.layout; - demand.source.source_type = transformed ? request.source_type : GGML_TYPE_COUNT; - demand.source.input_size = request.input_size; - demand.source.output_size = request.output_size; - result.weights.push_back(std::move(demand)); - } - return result; -} - CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, const CommandProgramBindings & bindings) { - return bind_command_program_weight_demands(build_command_program_weight_requests(commands), bindings); + return collect_weight_demands(commands, bindings, collect_graph_value_access(commands)); } bool debug_serial_command_execution_enabled() { diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index 42773cc1a2cd..7a35a01080a7 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -67,18 +67,6 @@ struct CommandProgramExecutionContext { CompileTraceExecution * compile_trace = nullptr; }; -// Requests contain command metadata only. Every occurrence remains available for validation before -// its deduplication group is emitted, since same-layout requests can differ in validity. -struct CommandProgramWeightRequest { - CommandBinding binding; - size_t group = 0; -}; - -struct CommandProgramWeightRequests { - std::vector requests; - size_t group_count = 0; -}; - struct CommandProgramWeightDemand { CommandBinding request; HostWeightSource source; @@ -92,10 +80,6 @@ struct CommandProgramWeightDemands { bool valid() const { return status.success(); } }; -CommandProgramWeightRequests build_command_program_weight_requests(const CommandProgram & commands); -CommandProgramWeightDemands bind_command_program_weight_demands(const CommandProgramWeightRequests & requests, - const CommandProgramBindings & bindings); - CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, const CommandProgramBindings & bindings); @@ -138,8 +122,7 @@ struct PreparedCommandProgram { std::vector initialization_commands; std::vector commands; std::vector host_staging; - std::shared_ptr> resident_host_weights = - std::make_shared>(); + std::vector resident_host_weights; std::vector program_constants; Status status; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; @@ -151,8 +134,6 @@ struct RecordedCommandGraph { hrx_graph_t graph = nullptr; hrx_graph_exec_t exec = nullptr; std::vector retained_buffers; - std::shared_ptr> resident_host_weights = - std::make_shared>(); std::vector retained_executables; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; size_t dispatch_count = 0; @@ -170,12 +151,6 @@ struct RecordedCommandGraph { bool valid() const { return status.success() && exec != nullptr; } }; -// Cached programs own their bindings across submissions; recorded handles are released before prepared leases. -struct CachedCommandProgramResources { - PreparedCommandProgram prepared; - RecordedCommandGraph recorded; -}; - struct RecordedCommandGraphExecutionResult { bool success = false; Status status; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 3673f1c7c11b..d192393596a4 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -130,10 +130,6 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { context_.prepared_programs.execute_with_result(execution_context, 0, lookup.program->command_shape_hash(), lookup.program->commands(), bindings); result.graph_replay_event = execution.graph_replay_event; - if (execution.execution_resources != nullptr) { - // A failed return can still follow queued writes. - lookup.program->mark_imported_writes(graph); - } if (!execution.success) { result.status.append(execution.status); if (result.status.success()) { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index 2c83f7d94998..af900bebc0e8 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -10,7 +10,6 @@ #include #include #include -#include #include #include #include @@ -412,7 +411,7 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur return result; } std::lock_guard entry_lock(prepared_entry_->mutex); - const auto & staging_buffers = prepared_entry_->resources->prepared.host_staging; + const auto & staging_buffers = prepared_entry_->prepared.host_staging; result.external_bindings.reserve(staging_buffers.size()); for (const HostStagingBuffer & staging : staging_buffers) { const GraphProgramExternalSlot * slot = find_external_slot(ValueId(staging.value)); @@ -432,26 +431,8 @@ GraphProgramMatch GraphProgram::match_host_staging_graph(const ggml_cgraph & cur Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const GraphProgramMatch & match) { Status status; external_slots_.clear(); - written_slots_.clear(); - std::unordered_set written_values; external_slot_by_value_.clear(); external_slots_.reserve(match.external_bindings.size()); - auto collect_access = [&](const std::vector & commands) { - for (const Command & command : commands) { - for (const CommandBinding & binding : command.bindings) { - const bool graph_value = binding.origin == CommandBindingOrigin::GraphValue; - const bool writes_only = binding.access == ResourceAccess::Write; - const bool reads_and_writes = binding.access == ResourceAccess::ReadWrite; - const bool writes_value = writes_only || reads_and_writes; - const bool graph_write = graph_value && writes_value; - if (graph_write) { - written_values.insert(binding.value.value); - } - } - } - }; - collect_access(commands_->initialization_commands); - collect_access(commands_->commands); for (const GraphProgramExternalBinding & binding : match.external_bindings) { GraphProgramExternalSlot slot; slot.value = binding.value; @@ -487,9 +468,6 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra } external_slot_by_value_[binding.value.value] = external_slots_.size(); external_slots_.push_back(slot); - if (written_values.count(slot.value.value) != 0) { - written_slots_.push_back(slot); - } } return status; } @@ -552,25 +530,6 @@ bool GraphProgram::mutable_sources_current() const { return true; } -void GraphProgram::mark_imported_writes(const ggml_cgraph & graph) const { - std::unordered_set marked_identities; - for (const GraphProgramExternalSlot & slot : written_slots_) { - Status status; - const ggml_tensor * tensor = resolve_external_slot(graph, slot, status); - if (tensor == nullptr) { - continue; - } - ggml_backend_hrx_imported_region region = {}; - if (!ggml_backend_hrx_get_retained_region(tensor, region)) { - continue; - } - const bool first_write = marked_identities.insert(region.identity).second; - if (first_write) { - ggml_backend_hrx_mark_retained_region_written(tensor); - } - } -} - GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_graph) const { GraphProgramMatch result; if (graph_ == nullptr) { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index d6d57945b6dd..54c41ca6873c 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -12,7 +12,6 @@ #include #include #include -#include #include #include @@ -75,7 +74,6 @@ class GraphProgram { bool can_use_prepared_fast_path(const ggml_cgraph & graph) const; bool mutable_sources_current() const; - void mark_imported_writes(const ggml_cgraph & graph) const; std::shared_ptr prepared_entry() const; void bind_prepared_graph(const ggml_cgraph & graph, const GraphProgramMatch & match, const CommandProgramBindings & bindings, @@ -102,7 +100,6 @@ class GraphProgram { size_t length; }; std::vector mutable_source_watches_; - std::vector written_slots_; std::unordered_map external_slot_by_value_; uint64_t fast_path_uid_ = 0; const ggml_tensor * const * fast_path_nodes_ = nullptr; diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 156dfe452203..73431ce93380 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -25,14 +25,6 @@ static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; static constexpr size_t kLargeHostUploadBytes = 1024 * 1024; static constexpr size_t kQ6PackScratchLimitBytes = 64 * 1024 * 1024; -static bool mapped_download_staging_enabled() { - static const bool enabled = [] { - const char * mode = std::getenv("GGML_HRX_WEIGHT_DOWNLOAD_STAGING"); - return mode == nullptr || mode[0] == '\0' || std::strcmp(mode, "mapped") == 0; - }(); - return enabled; -} - struct SymmetricI4Block { std::array scales = {}; std::array, 8> payloads = {}; @@ -1261,13 +1253,6 @@ Status allocate_mapped_host_staging_buffer(hrx_device_t device, size_t size, hrx return status; } -HostTransferManager::~HostTransferManager() { - std::lock_guard lock(mutex_); - if (mapped_download_buffer_ != nullptr) { - hrx_buffer_release(mapped_download_buffer_); - } -} - Status HostTransferManager::upload_synchronous(hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, @@ -1371,55 +1356,6 @@ Status HostTransferManager::download_synchronous(hrx_stream_t stream, return status; } -Status HostTransferManager::download_to_mapped_synchronous(hrx_device_t device, - hrx_stream_t stream, - hrx_buffer_t source, - size_t offset, - size_t size, - HostMappedDownloadLease & lease) { - Status status; - lease = {}; - if (device == nullptr || stream == nullptr || source == nullptr || size == 0) { - status.log("invalid mapped HRX host download"); - return status; - } - lease.lock = std::unique_lock(mutex_); - if (mapped_download_capacity_ < size) { - size_t capacity = std::max(mapped_download_capacity_, 1); - while (capacity < size && capacity <= std::numeric_limits::max() / 2) { - capacity *= 2; - } - if (capacity < size) { - capacity = size; - } - hrx_buffer_t new_buffer = nullptr; - void * new_data = nullptr; - status = allocate_mapped_host_staging_buffer(device, capacity, new_buffer, new_data); - if (!status.success()) { - return status; - } - if (mapped_download_buffer_ != nullptr) { - hrx_buffer_release(mapped_download_buffer_); - } - mapped_download_buffer_ = new_buffer; - mapped_download_data_ = new_data; - mapped_download_capacity_ = capacity; - } - if (ErrorResult error = - take_status(hrx_stream_copy_buffer(stream, source, offset, mapped_download_buffer_, 0, size))) { - status.log("mapped host weight download failed: %s", error->c_str()); - return status; - } - if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { - status.log("synchronize mapped host weight download failed: %s", error->c_str()); - return status; - } - ++stats_.downloads; - stats_.download_bytes += size; - lease.host_data = mapped_download_data_; - return status; -} - HostTransferStats HostTransferManager::stats() const { std::lock_guard lock(mutex_); return stats_; @@ -1430,29 +1366,16 @@ void HostTransferManager::clear() { stats_ = {}; } -struct HostWeightAccounting { - std::mutex mutex; - size_t allocation_count = 0; - size_t resident_bytes = 0; - size_t peak_resident_bytes = 0; -}; - struct HostWeightLease::Entry { ~Entry() { if (buffer != nullptr) { hrx_buffer_release(buffer); } - if (accounting != nullptr) { - std::lock_guard lock(accounting->mutex); - --accounting->allocation_count; - accounting->resident_bytes -= length; - } } hrx_buffer_t buffer = nullptr; size_t length = 0; std::string layout; - std::shared_ptr accounting; }; HostWeightLease::HostWeightLease(std::shared_ptr entry) : entry_(std::move(entry)) {} @@ -1474,8 +1397,6 @@ const std::string & HostWeightLease::layout() const { return entry_ != nullptr ? entry_->layout : empty; } -HostWeightCache::HostWeightCache() : accounting_(std::make_shared()) {} - HostWeightCache::~HostWeightCache() { clear(); } @@ -1510,10 +1431,14 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device result.status.log("host weight residency requires an HRX stream"); return result; } - const bool has_host_source = source.host_data != nullptr; - const bool has_device_source = source.device_buffer != nullptr; - if (has_host_source == has_device_source || source.identity == 0 || source.generation == 0 || source.length == 0 || - source.offset > source.capacity || source.length > source.capacity - source.offset) { + const bool has_host_source = source.host_data != nullptr; + const bool has_identity = source.identity != 0; + const bool has_generation = source.generation != 0; + const bool has_length = source.length != 0; + const bool offset_fits = source.offset <= source.capacity; + const bool length_fits = offset_fits && source.length <= source.capacity - source.offset; + const bool valid_source = has_host_source && has_identity && has_generation && has_length && length_fits; + if (!valid_source) { result.status.log("invalid host weight source"); return result; } @@ -1551,35 +1476,6 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device --fail_cold_after_; } - HostWeightSource materialization_source = source; - HostMappedDownloadLease mapped_canonical; - std::vector canonical; - if (has_device_source) { - if (source.layout == kNativeWeightLayout) { - result.status.log("native device weights do not require host materialization"); - return result; - } - if (!mapped_download_staging_enabled()) { - canonical.resize(source.length); - } - void * canonical_data = canonical.data(); - if (mapped_download_staging_enabled()) { - result.status = transfers.download_to_mapped_synchronous( - device, stream, source.device_buffer, source.offset, source.length, mapped_canonical); - canonical_data = mapped_canonical.host_data; - } else { - result.status = transfers.download_synchronous(stream, source.device_buffer, source.offset, canonical_data, - source.length); - } - if (!result.status.success()) { - return result; - } - materialization_source.host_data = canonical_data; - materialization_source.device_buffer = nullptr; - materialization_source.capacity = source.length; - materialization_source.offset = 0; - } - const bool q6_packed = source.layout == kQ6KPackedK256Row64ScaleRowLayout; const bool scratch_fits = materialized_length <= kQ6PackScratchLimitBytes; auto scratch = q6_pack_scratch_.find(materialized_length); @@ -1598,10 +1494,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device // Budget 64 MiB of host scratch across first-use materializations. The Q6 packer // writes every output byte, and the cache lock covers packing and upload. std::vector & output = use_scratch ? scratch->second : transformed; - result.status = materialize_weight(materialization_source, upload_data, upload_size, output); - if (mapped_canonical.lock.owns_lock()) { - mapped_canonical.lock.unlock(); - } + result.status = materialize_weight(source, upload_data, upload_size, output); if (!result.status.success()) { return result; } @@ -1618,19 +1511,15 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device if (!result.status.success()) { return result; } - entry->accounting = accounting_; - { - std::lock_guard accounting_lock(accounting_->mutex); - ++accounting_->allocation_count; - accounting_->resident_bytes += materialized_length; - accounting_->peak_resident_bytes = std::max(accounting_->peak_resident_bytes, accounting_->resident_bytes); - } result.status = transfers.upload_synchronous(stream, upload_data, entry->buffer, 0, upload_size); if (!result.status.success()) { return result; } entries_.emplace(key, entry); + ++stats_.allocation_count; + stats_.resident_bytes += materialized_length; + stats_.peak_resident_bytes = std::max(stats_.peak_resident_bytes, stats_.resident_bytes); report_new_representation(key, *entry); ++stats_.misses; stats_.materialized_bytes += materialized_length; @@ -1640,12 +1529,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device HostWeightCacheStats HostWeightCache::stats() const { std::lock_guard lock(mutex_); - std::lock_guard accounting_lock(accounting_->mutex); - HostWeightCacheStats result = stats_; - result.allocation_count = accounting_->allocation_count; - result.resident_bytes = accounting_->resident_bytes; - result.peak_resident_bytes = accounting_->peak_resident_bytes; - return result; + return stats_; } void HostWeightCache::report_new_representation(const SourceKey & key, const HostWeightLease::Entry & entry) { @@ -1687,6 +1571,8 @@ void HostWeightCache::fail_cold_acquire_after_for_test(size_t successful_acquire void HostWeightCache::clear() { std::lock_guard lock(mutex_); entries_.clear(); + stats_.allocation_count = 0; + stats_.resident_bytes = 0; q6_pack_scratch_.clear(); q6_pack_scratch_bytes_ = 0; } diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 7cd4197228fd..87841e4baf23 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -26,15 +26,8 @@ struct HostTransferStats { size_t download_bytes = 0; }; -struct HostMappedDownloadLease { - std::unique_lock lock; - void * host_data = nullptr; -}; - class HostTransferManager { public: - ~HostTransferManager(); - Status upload_synchronous(hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, @@ -50,27 +43,16 @@ class HostTransferManager { size_t offset, void * host_destination, size_t size); - Status download_to_mapped_synchronous(hrx_device_t device, - hrx_stream_t stream, - hrx_buffer_t source, - size_t offset, - size_t size, - HostMappedDownloadLease & lease); - HostTransferStats stats() const; void clear(); private: mutable std::mutex mutex_; HostTransferStats stats_; - hrx_buffer_t mapped_download_buffer_ = nullptr; - void * mapped_download_data_ = nullptr; - size_t mapped_download_capacity_ = 0; }; struct HostWeightSource { const void * host_data = nullptr; - hrx_buffer_t device_buffer = nullptr; uint64_t identity = 0; uint64_t generation = 0; size_t capacity = 0; @@ -83,6 +65,7 @@ struct HostWeightSource { int64_t output_size = 0; }; +// Residency counts cache-owned representations, not buffers retained after an explicit clear. struct HostWeightCacheStats { uint64_t hits = 0; uint64_t misses = 0; @@ -93,8 +76,6 @@ struct HostWeightCacheStats { uint64_t duplicate_warnings = 0; }; -struct HostWeightAccounting; - class HostWeightLease { public: HostWeightLease() = default; @@ -122,7 +103,7 @@ struct HostWeightAcquireResult { // Owns each immutable source representation until explicit clearing or backend teardown. class HostWeightCache { public: - HostWeightCache(); + HostWeightCache() = default; ~HostWeightCache(); HostWeightCache(const HostWeightCache &) = delete; @@ -134,7 +115,7 @@ class HostWeightCache { const HostWeightSource & source); HostWeightCacheStats stats() const; void fail_cold_acquire_after_for_test(size_t successful_acquires); - // Erases lookup ownership; leases and cumulative diagnostics outlive map entries. + // Clears cache residency; existing leases remain valid and cumulative diagnostics are preserved. void clear(); private: @@ -171,7 +152,6 @@ class HostWeightCache { std::unordered_map> q6_pack_scratch_; size_t q6_pack_scratch_bytes_ = 0; HostWeightCacheStats stats_; - std::shared_ptr accounting_; size_t fail_cold_after_ = SIZE_MAX; }; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index 31fb48bdeb54..5ab1fcc193da 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -100,15 +100,13 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ if (graph_uid == 0 || !commands.valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; result.graph_replay_ineligible_reason = "uncached_graph"; - auto resources = std::make_shared(); - resources->prepared = prepare_command_program(context, commands, bindings); - if (!resources->prepared.valid()) { - result.status.append(resources->prepared.status); + auto prepared = prepare_command_program(context, commands, bindings); + if (!prepared.valid()) { + result.status.append(prepared.status); return result; } - result.execution_resources = resources; // HRX command buffers retain queued buffers and executables even after this uncached owner is released. - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, prepared); if (!result.success) { result.status.log("execute uncached HRX command program failed"); } @@ -145,84 +143,44 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ const bool replay_disabled = graph_replay_disabled_by_environment(); const bool serial_execution = debug_serial_command_execution_enabled(); const bool execute_directly = replay_disabled || serial_execution; - if (entry->resources != nullptr && entry->resources->prepared.valid()) { - auto & resources = *entry->resources; - result.execution_resources = entry->resources; - record_hit(); - if (execute_directly) { - result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : - "disabled_by_environment"; - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources.prepared); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } + if (!entry->ready) { + auto prepared = prepare_command_program(context, commands, bindings); + if (!prepared.valid()) { + result.status.append(prepared.status); return result; } + entry->prepared = std::move(prepared); + entry->ready = true; + record_build(); + } else { + record_hit(); + } + + if (execute_directly) { + result.graph_replay_event = HrxGraphReplayEvent::Disabled; + result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : + "disabled_by_environment"; + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->prepared); + } else { const char * rebuild_setting = std::getenv("GGML_HRX_REBUILD_GRAPH_REPLAY"); const bool rebuild_requested = rebuild_setting != nullptr && rebuild_setting[0] != '\0'; const bool rebuild_enabled = rebuild_requested && rebuild_setting[0] != '0'; if (rebuild_enabled) { - resources.recorded = {}; + entry->recorded = {}; } - const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, commands, bindings, resources.prepared, resources.recorded); + const auto replay = bind_and_launch_recorded_command_graph( + context, commands, bindings, entry->prepared, entry->recorded); apply_graph_replay_result(result, replay); - if (replay.success) { - result.success = true; - return result; - } - if (!graph_replay_should_fallback(replay.event)) { + result.success = replay.success; + const bool fallback = !replay.success && graph_replay_should_fallback(replay.event); + if (fallback) { + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->prepared); + } else if (!replay.success) { result.status.append(replay.status); - if (result.status.success()) { - result.status.log("execute cached HRX graph replay failed"); - } - return result; - } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources.prepared); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } - return result; - } - - auto resources = std::make_shared(); - resources->prepared = prepare_command_program(context, commands, bindings); - if (!resources->prepared.valid()) { - result.status.append(resources->prepared.status); - return result; - } - entry->resources = resources; - result.execution_resources = resources; - record_build(); - - if (execute_directly) { - result.graph_replay_event = HrxGraphReplayEvent::Disabled; - result.graph_replay_ineligible_reason = serial_execution ? "debug_serial_execution" : - "disabled_by_environment"; - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); - if (!result.success) { - result.status.log("execute prepared HRX command program failed"); } - return result; - } - - const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, commands, bindings, resources->prepared, resources->recorded); - apply_graph_replay_result(result, replay); - if (replay.success) { - result.success = true; - return result; - } - if (!graph_replay_should_fallback(replay.event)) { - result.status.append(replay.status); - if (result.status.success()) { - result.status.log("execute prepared HRX graph replay failed"); - } - return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, resources->prepared); - if (!result.success) { + const bool missing_error = !result.success && result.status.success(); + if (missing_error) { result.status.log("execute prepared HRX command program failed"); } return result; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index a5f713301d8a..d18ef44375e9 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -24,7 +24,10 @@ struct PreparedCommandProgramCacheStats { struct PreparedCommandProgramCacheEntry { std::mutex mutex; - std::shared_ptr resources; + // Release recorded handles before the prepared bindings they reference. + PreparedCommandProgram prepared; + RecordedCommandGraph recorded; + bool ready = false; }; struct PreparedCommandProgramCacheExecutionResult { @@ -37,7 +40,6 @@ struct PreparedCommandProgramCacheExecutionResult { uint64_t graph_replay_total_ns = 0; size_t graph_replay_dispatches = 0; bool graph_replay_transient_allocation_changed = false; - std::shared_ptr execution_resources; }; class PreparedCommandProgramCache { diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 4247d19489d9..38beffc5e300 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -11768,12 +11768,14 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { lookup.program->commands(), bindings); REQUIRE(execution.success); ggml_backend_synchronize(backend); - const auto & prepared = execution.execution_resources->prepared; + const auto entry = state->prepared_programs.get_or_create( + execution_context, lookup.program->uid(), lookup.program->command_shape_hash(), bindings); + const auto & prepared = entry->prepared; REQUIRE(std::none_of(prepared.host_staging.begin(), prepared.host_staging.end(), [&](const auto & staging) { return staging.value == unused_value->id.value; })); REQUIRE(prepared.host_staging.size() == 2); - REQUIRE(prepared.resident_host_weights->size() == 1); + REQUIRE(prepared.resident_host_weights.size() == 1); REQUIRE(state->host_weights.stats().resident_bytes == weight_bytes); REQUIRE(state->host_weights.stats().misses == 1); // The vector matmul declares both host input and output as reads. @@ -11795,7 +11797,7 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { auto prepared = ggml::hrx::prepare_command_program(execution_context, initialization_only, bindings); REQUIRE(prepared.valid()); REQUIRE(prepared.host_staging.size() == 2); - REQUIRE(prepared.resident_host_weights->size() == 1); + REQUIRE(prepared.resident_host_weights.size() == 1); REQUIRE(state->host_weights.stats().misses == 1); } ggml_backend_buffer_free(cpu_buffer); @@ -12303,7 +12305,7 @@ static void run_disjoint_transient_plan_packing_checks() { ggml_free(ctx); } -static void run_command_weight_request_checks() { +static void run_command_weight_demand_checks() { using namespace ggml::hrx; std::array source = {}; CommandProgramBinding binding; @@ -12359,7 +12361,6 @@ static void run_command_weight_request_checks() { CommandProgram commands; commands.initialization_commands.push_back(init); commands.commands.push_back(main); - const auto requests = build_command_program_weight_requests(commands); std::vector current; for (int32_t value = 1; value <= 5; ++value) { auto next = binding; @@ -12368,7 +12369,7 @@ static void run_command_weight_request_checks() { } const auto bindings = CommandProgramBindings::from_bindings(current); REQUIRE(bindings.valid()); - const auto demands = bind_command_program_weight_demands(requests, bindings); + const auto demands = collect_command_program_weight_demands(commands, bindings); REQUIRE(demands.valid()); REQUIRE(demands.weights.size() == 4); REQUIRE(demands.weights[0].request.name == "first-native"); @@ -12381,14 +12382,7 @@ static void run_command_weight_request_checks() { REQUIRE(demands.weights[1].source.length == 64); REQUIRE(demands.weights[3].source.length == 32); REQUIRE(demands.weights[3].source.input_size == 4); - const auto collected = collect_command_program_weight_demands(commands, bindings); - REQUIRE(collected.valid()); - REQUIRE(collected.weights.size() == demands.weights.size()); - REQUIRE(collected.weights[2].request.value == ValueId(2)); - - // Reuse the immutable plan with different live source metadata and allocation eligibility. - commands.initialization_commands.clear(); - commands.commands.clear(); + // Collect demand against the current source metadata and allocation eligibility. for (auto & next : current) { next.buffer = nullptr; next.host_data = source.data() + 32; @@ -12399,8 +12393,8 @@ static void run_command_weight_request_checks() { next.source_offset = 32; next.source_capacity = source.size() - 16; } - const auto rebound = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings(current)); + const auto rebound = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings(current)); REQUIRE(rebound.valid()); REQUIRE(rebound.weights.size() == 4); for (const auto & demand : rebound.weights) { @@ -12420,35 +12414,34 @@ static void run_command_weight_request_checks() { next.weight = false; next.source_host_data = nullptr; } - const auto nonweights = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings(current)); + const auto nonweights = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings(current)); REQUIRE(nonweights.valid()); REQUIRE(nonweights.weights.empty()); CommandProgram native_program; native_program.initialization_commands.push_back(Command{}); native_program.initialization_commands.back().bindings.push_back(native); - const auto native_requests = build_command_program_weight_requests(native_program); auto uploaded = binding; uploaded.source_host_data = nullptr; - const auto uploaded_demand = bind_command_program_weight_demands( - native_requests, CommandProgramBindings::from_bindings({ uploaded })); + const auto uploaded_demand = collect_command_program_weight_demands( + native_program, CommandProgramBindings::from_bindings({ uploaded })); REQUIRE(uploaded_demand.valid()); REQUIRE(uploaded_demand.weights.empty()); - const auto restored = bind_command_program_weight_demands(native_requests, bindings); + const auto restored = collect_command_program_weight_demands(native_program, bindings); REQUIRE(restored.valid()); REQUIRE(restored.weights.size() == 1); REQUIRE(!restored.weights[0].requires_materialization); auto later_binding = binding; later_binding.identity = 99; - const auto first_binding = bind_command_program_weight_demands( - native_requests, CommandProgramBindings::from_bindings({ binding, later_binding })); + const auto first_binding = collect_command_program_weight_demands( + native_program, CommandProgramBindings::from_bindings({ binding, later_binding })); REQUIRE(first_binding.valid()); REQUIRE(first_binding.weights.size() == 1); REQUIRE(first_binding.weights[0].source.identity == binding.identity); } -static void run_command_weight_request_validation_checks() { +static void run_command_weight_demand_validation_checks() { using namespace ggml::hrx; std::array source = {}; CommandProgramBinding binding; @@ -12476,9 +12469,8 @@ static void run_command_weight_request_validation_checks() { Command main; main.bindings.push_back(invalid_first ? valid : invalid_offset); commands.commands.push_back(main); - const auto requests = build_command_program_weight_requests(commands); - const auto demands = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings({ binding })); + const auto demands = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ binding })); REQUIRE(!demands.valid()); REQUIRE(demands.status.errors().size() == 1); REQUIRE(demands.status.errors()[0].find("invalid transformed layout request") != std::string::npos); @@ -12487,8 +12479,8 @@ static void run_command_weight_request_validation_checks() { REQUIRE(demands.weights[0].request.offset == 0); auto missing = binding; missing.source_host_data = nullptr; - const auto missing_source = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings({ missing })); + const auto missing_source = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ missing })); REQUIRE(!missing_source.valid()); REQUIRE(missing_source.weights.empty()); REQUIRE(missing_source.status.errors().size() == 2); @@ -12496,11 +12488,11 @@ static void run_command_weight_request_validation_checks() { REQUIRE(error.find("no retained source") != std::string::npos); } missing.weight = false; - const auto nonweight = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings({ missing })); + const auto nonweight = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ missing })); REQUIRE(nonweight.valid()); REQUIRE(nonweight.weights.empty()); - const auto unbound = bind_command_program_weight_demands(requests, CommandProgramBindings{}); + const auto unbound = collect_command_program_weight_demands(commands, CommandProgramBindings{}); REQUIRE(unbound.valid()); REQUIRE(unbound.weights.empty()); } @@ -12527,15 +12519,14 @@ static void run_command_weight_request_validation_checks() { Command init; init.bindings.push_back(valid); commands.initialization_commands.push_back(init); - const auto requests = build_command_program_weight_requests(commands); auto short_binding = binding; short_binding.length = valid.source_length - 1; - const auto too_short = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings({ short_binding })); + const auto too_short = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ short_binding })); REQUIRE(!too_short.valid()); REQUIRE(too_short.weights.empty()); - const auto restored = bind_command_program_weight_demands( - requests, CommandProgramBindings::from_bindings({ binding })); + const auto restored = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ binding })); REQUIRE(restored.valid()); REQUIRE(restored.weights.size() == 1); } @@ -14667,8 +14658,8 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("multiple_transient_plan", [] { run_multiple_transient_plan_checks(); }); suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); - suite.host_case("command_weight_requests", [] { run_command_weight_request_checks(); }); - suite.host_case("command_weight_request_validation", [] { run_command_weight_request_validation_checks(); }); + suite.host_case("command_weight_demands", [] { run_command_weight_demand_checks(); }); + suite.host_case("command_weight_demand_validation", [] { run_command_weight_demand_validation_checks(); }); suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); suite.host_case("graph_compute_mask", [] { run_graph_compute_mask_checks(); }); suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index c136808c8845..185bd1b213bf 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -266,9 +266,6 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { REQUIRE(first_region.identity == second_region.identity); REQUIRE(first_region.offset == 0); REQUIRE(second_region.offset == 32 * sizeof(uint32_t)); - REQUIRE(first_region.type == GGML_TYPE_I32); - REQUIRE(first_region.ne[0] == 32); - REQUIRE(first_region.nb[0] == sizeof(uint32_t)); ggml_tensor * view = ggml_view_1d(ggml, second, 8, 4 * sizeof(uint32_t)); REQUIRE(ggml_backend_view_init(view) == GGML_STATUS_SUCCESS); @@ -277,7 +274,6 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { REQUIRE(view_region.identity == second_region.identity); REQUIRE(view_region.offset == second_region.offset + 4 * sizeof(uint32_t)); REQUIRE(view_region.length == 8 * sizeof(uint32_t)); - REQUIRE(view_region.ne[0] == 8); ggml::hrx::ValueBufferBinding binding; REQUIRE(ggml_backend_hrx_resolve_value_buffer(second, binding)); @@ -592,25 +588,23 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { ggml::hrx::RecordedCommandGraph recorded; hrx_buffer_retain(first.lease.buffer()); recorded.retained_buffers.push_back(first.lease.buffer()); - recorded.resident_host_weights = - std::make_shared>(std::initializer_list{ first.lease }); weights.clear(); weight_stats = weights.stats(); REQUIRE(weight_stats.hits == 1); REQUIRE(weight_stats.misses == 3); - REQUIRE(weight_stats.allocation_count == 3); - REQUIRE(weight_stats.resident_bytes == 96); + REQUIRE(weight_stats.allocation_count == 0); + REQUIRE(weight_stats.resident_bytes == 0); + REQUIRE(weight_stats.peak_resident_bytes == 96); first.lease = {}; second.lease = {}; slice.lease = {}; next_generation.lease = {}; - weight_stats = weights.stats(); - REQUIRE(weight_stats.allocation_count == 1); - REQUIRE(weight_stats.resident_bytes == 32); + // Recording retention keeps the buffer valid independently of cache residency and weight leases. + std::array retained_bytes = {}; + require_hrx_status(hrx_synchronous_d2h(context->device->device, recorded.retained_buffers.front(), 0, + retained_bytes.data(), retained_bytes.size())); + REQUIRE(retained_bytes == first_bytes); recorded = {}; - weight_stats = weights.stats(); - REQUIRE(weight_stats.allocation_count == 0); - REQUIRE(weight_stats.resident_bytes == 0); source.layout = ggml::hrx::kNativeWeightLayout; source.generation = 3; diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index ebe425bf7e1a..db2830d1f001 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -2863,9 +2863,6 @@ static void run_imported_native_weight_case() { ggml_backend_hrx_imported_region region = {}; REQUIRE(ggml_backend_hrx_get_imported_region(weight, region)); - REQUIRE(region.type == GGML_TYPE_F32); - REQUIRE(region.ne[0] == kInputSize); - REQUIRE(region.ne[1] == kOutputSize); REQUIRE(region.offset == 0); ggml_backend_hrx_imported_region alias_region = {}; REQUIRE(ggml_backend_hrx_get_imported_region(weight_alias, alias_region)); @@ -3315,16 +3312,30 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { &ggml::hrx::get_qwen_kernel_corpus(), &state->kernel_executables, &state->transient_arena, &state->host_transfers, &state->host_weights, &state->device->host_buffers, &state->graph_replay_state, }; - auto execution = state->prepared_programs.execute_with_result( - execution_context, uncached ? 0 : packed_lookup.program->uid(), packed_lookup.program->command_shape_hash(), - packed_lookup.program->commands(), packed_bindings); - REQUIRE(execution.success); - REQUIRE(!execution.execution_resources->prepared.host_staging.empty()); - REQUIRE(execution.execution_resources->recorded.valid() == !uncached); + std::shared_ptr entry; + std::unique_ptr uncached_prepared; + ggml::hrx::PreparedCommandProgram * prepared = nullptr; + if (uncached) { + uncached_prepared = std::make_unique( + ggml::hrx::prepare_command_program(execution_context, packed_lookup.program->commands(), packed_bindings)); + REQUIRE(uncached_prepared->valid()); + REQUIRE(ggml::hrx::bind_and_execute_prepared_command_program( + execution_context, packed_lookup.program->commands(), packed_bindings, *uncached_prepared)); + prepared = uncached_prepared.get(); + } else { + entry = state->prepared_programs.get_or_create( + execution_context, packed_lookup.program->uid(), packed_lookup.program->command_shape_hash(), packed_bindings); + const auto execution = state->prepared_programs.execute_with_result( + execution_context, packed_lookup.program->commands(), packed_bindings, entry); + REQUIRE(execution.success); + REQUIRE(entry->recorded.valid()); + prepared = &entry->prepared; + } + REQUIRE(!prepared->host_staging.empty()); ggml_backend_synchronize(backend); ggml_backend_tensor_set(hrx_output, poisoned.data(), 0, poisoned.size() * sizeof(float)); // The invalid final command fails after the real weight reader is queued, before the completion boundary. - execution.execution_resources->prepared.commands.emplace_back(); + prepared->commands.emplace_back(); hrx_semaphore_t gate = nullptr; REQUIRE(!ggml::hrx::take_status(hrx_semaphore_create(state->device->device, 0, &gate))); REQUIRE(!ggml::hrx::take_status(hrx_stream_wait_on(state->stream, { gate, 1 }))); @@ -3341,10 +3352,11 @@ static void run_failed_submission_weight_lifetime_case(bool uncached = false) { } }); const bool execution_failed = !ggml::hrx::execute_prepared_command_program( - execution_context, execution.execution_resources->prepared); - std::weak_ptr sole_resource = execution.execution_resources; + execution_context, *prepared); + std::weak_ptr sole_resource = entry; const bool flush_ok = !ggml::hrx::take_status(hrx_stream_flush(state->stream)); - execution.execution_resources.reset(); + entry.reset(); + uncached_prepared.reset(); if (uncached) { state->graph_programs.clear(); state->kernel_executables.clear(); @@ -3611,7 +3623,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, auto prepared = ggml::hrx::prepare_command_program( execution_context, lookup.program->commands(), bindings); REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights->size() == expected_materializations); + REQUIRE(prepared.resident_host_weights.size() == expected_materializations); const bool staged_weight = std::any_of(prepared.host_staging.begin(), prepared.host_staging.end(), [&](const auto & staging) { return staging.value == weight_value->id.value; }); REQUIRE(!staged_weight); @@ -3846,8 +3858,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, REQUIRE(after_rebind.weight_materializations == before_rebind.weight_materializations); { - // Exercise production preparation and recording: physical residency must remain counted after both - // the graph-program owner and weight lookup entries are gone, for as long as recording retains it. + // A recording retains its buffers even after prepared bindings and cache ownership are gone. auto full_match = lookup.program->match_trusted_graph(*hrx_decode, true); REQUIRE(full_match.valid()); const ggml::hrx::CommandProgram & commands = lookup.program->commands(); @@ -3871,7 +3882,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, }; auto prepared = ggml::hrx::prepare_command_program(execution_context, commands, bindings); REQUIRE(prepared.valid()); - REQUIRE(prepared.resident_host_weights->size() == 1); + REQUIRE(prepared.resident_host_weights.size() == 1); ggml::hrx::RecordedCommandGraph recorded; const auto launch = ggml::hrx::bind_and_launch_recorded_command_graph( execution_context, commands, bindings, prepared, recorded); @@ -3883,7 +3894,11 @@ static void run_imported_packed_weight_case(ggml_type weight_type, backend_state->prepared_programs.clear(); lookup = {}; backend_state->host_weights.clear(); - REQUIRE(backend_state->host_weights.stats().allocation_count == 1); + REQUIRE(backend_state->host_weights.stats().allocation_count == 0); + REQUIRE(!ggml::hrx::take_status(hrx_graph_exec_launch(recorded.exec, backend_state->stream))); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_decode_output), get_f32_tensor(cpu_backend, cpu_decode_output), + 2.0e-2f, 3.0e-2f); recorded = {}; REQUIRE(backend_state->host_weights.stats().allocation_count == 0); } From 1ee5b60baf325764ef6d095a5cd828e9cb231d1a Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Mon, 2 Nov 2026 23:21:34 +0000 Subject: [PATCH 13/18] tests: align checks with persistent weight ownership Assisted-by: Codex --- tests/hrx-buffer-test.cpp | 5 ----- tests/hrx-ops-test.cpp | 8 ++++++-- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index 185bd1b213bf..1c152f2627bc 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -390,13 +390,8 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { REQUIRE(coherent_graph != nullptr); ggml_build_forward_expand(coherent_graph, coherent_copy); coherent_graph->uid = 950005; - ggml_backend_hrx_imported_region coherent_region = {}; - REQUIRE(ggml_backend_hrx_get_imported_region(coherent_output, coherent_region)); - const uint64_t coherent_generation = coherent_region.generation; REQUIRE(ggml_backend_graph_compute(backend, coherent_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); - REQUIRE(ggml_backend_hrx_get_imported_region(coherent_output, coherent_region)); - REQUIRE(coherent_region.generation != coherent_generation); REQUIRE(*reinterpret_cast(subregion) == coherent_values[0]); ggml_backend_buffer_free(coherent_input_buffer); ggml_backend_buffer_free(cpu_buffer); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index db2830d1f001..8c1d64db7aa0 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3889,16 +3889,20 @@ static void run_imported_packed_weight_case(ggml_type weight_type, REQUIRE(launch.success); REQUIRE(recorded.valid()); ggml_backend_synchronize(hrx_backend); + const auto expected = get_f32_tensor(hrx_backend, hrx_decode_output); + require_close(expected, get_f32_tensor(cpu_backend, cpu_decode_output), 1.0f, 3.0e-2f); + const std::vector poisoned(expected.size(), 123.0f); + REQUIRE(expected != poisoned); prepared = {}; backend_state->graph_programs.clear(); backend_state->prepared_programs.clear(); lookup = {}; backend_state->host_weights.clear(); REQUIRE(backend_state->host_weights.stats().allocation_count == 0); + ggml_backend_tensor_set(hrx_decode_output, poisoned.data(), 0, poisoned.size() * sizeof(float)); REQUIRE(!ggml::hrx::take_status(hrx_graph_exec_launch(recorded.exec, backend_state->stream))); ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_decode_output), get_f32_tensor(cpu_backend, cpu_decode_output), - 2.0e-2f, 3.0e-2f); + REQUIRE(get_f32_tensor(hrx_backend, hrx_decode_output) == expected); recorded = {}; REQUIRE(backend_state->host_weights.stats().allocation_count == 0); } From 044d20046a5838ed26e075148464c4b0e5290733 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Tue, 3 Nov 2026 03:12:38 +0000 Subject: [PATCH 14/18] test(ggml-hrx): restore packed Q4 matrix coverage Run the existing gated and binary Q4 cases with imported weight sources, keeping native fallback coverage. Condense the weight residency diagnostics as requested. Refs: lemonade-ws-tek.33 Assisted-by: Codex --- ggml/src/ggml-hrx/README.md | 14 ++-------- tests/hrx-ops-test.cpp | 55 ++++++++++++++----------------------- 2 files changed, 23 insertions(+), 46 deletions(-) diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index 988f94c3b836..45b13ac6bc6b 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -2,19 +2,9 @@ ## Weight residency diagnostics -Weight residency supports standard GGUF inference with immutable model weights. Alternative layouts require an HRX `buffer_from_host_ptr` import that does not alias a known HRX host allocation. Other weights remain native-only: host weights reuse a cached native copy, and directly bindable device weights reuse their existing allocation. Source availability for native caching does not grant permission to select a transformed layout. +HRX materializes imported host weights in the layout requested by the selected kernel and caches each result for the backend's lifetime. If one weight is materialized in two layouts, HRX warns about the duplicate allocations and their memory cost. -Set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v` to emit one `ggml_hrx: weight_stats` line when each HRX backend is released. The `-v` flag enables the backend log callback in `llama-bench`. The line reports successful weight materializations, cache acquisition reuse, uploaded weight bytes, cached allocation count and bytes, peak cached bytes, and all host transfer uploads. `ggml_backend_hrx_get_cache_stats` exposes the same weight counters while a backend is running. Prepared program hits do not count as weight cache acquisitions, so unchanged replay can reuse a weight without increasing the `reused` counter. - -Resident weight bytes count physical allocations owned by the weight cache. An explicit clear resets the resident count and bytes while preserving cumulative materializations, reuse, warnings and peak bytes. Existing prepared leases and recorded buffer handles remain valid after a clear, but are no longer counted as cache residency. Imported mappings and directly bound coherent host buffers are source storage, not additional weight allocations. Native and losslessly packed layouts can both be materialized from their retained bytes; numerical conversion layouts retain their existing conversion behavior. - -Each immutable source region is materialized only when a command first requests its layout. Native residency copies the identified source byte range verbatim, so tensor type, shape, and strides do not create separate native copies. Transformed representations additionally distinguish layout, source type, transformation dimensions, and source/output byte lengths. Graph metadata still governs how those bytes are interpreted. The weight cache keeps every distinct representation until explicit cache clearing or backend teardown. Prefill/decode transitions, empty graphs, and other weights do not evict it. Prepared programs hold allocation leases; recorded graphs retain the HRX buffer handles they use. - -A graph instance remembers its validated prepared-cache entry. Its first execution resolves all physical bindings; repeated executions with the same graph UID and node array bind only mutable host staging. Structurally equivalent graphs can share command programs without sharing unchecked bindings: each new instance goes through the full binding lookup before using the fast path. Mutable source rebinding also returns to that lookup. Both execution paths use the same prepared-cache entry, lock, and replay implementation; warm executions do not rebuild per-weight demand. On a preparation miss, one command-access analysis covers staging and weight demand. Every layout request is validated against the current source before equivalent demands are deduplicated. - -Inserting a new physical representation logs a warning when another layout of the same source region and generation is already cached. The warning groups exact source byte ranges; it does not detect partial overlap or infer a whole tensor from a transformed subrange. It names the layouts, the new allocation's byte cost, and their total resident bytes. Aliases reuse the same representation, and warm executions do not repeat the warning. Keeping both layouts can increase steady-state device memory, while avoiding repeated conversion and upload. - -Keep borrowed weight bytes unchanged and alive while HRX can use them. A foreign CPU host buffer and its storage must stay alive for the lifetime of the HRX backend that caches it. Its identity depends on its wrapper and base addresses; destroying and recreating a buffer at those addresses while that backend lives can reuse old residency. Training and in-place weight updates are outside this path. Ordinary mutable inputs and outputs retain their host staging, writeback and synchronization behavior. +To log residency totals, set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v`. HRX prints a `ggml_hrx: weight_stats` line when the backend is released, with materialization and reuse counts, cached allocations and bytes, peak cached bytes, and upload totals. ## Compilation diagnostics diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 8c1d64db7aa0..93470c8ecc5c 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -8265,7 +8265,8 @@ static void register_dense_matmul_cases(Suite & suite) { [op, outputs] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 4, outputs, op); + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, outputs, op, + kQwenHiddenSize, false, true); }); } for (const int64_t tokens : { 256, 512 }) { @@ -8275,7 +8276,8 @@ static void register_dense_matmul_cases(Suite & suite) { [op] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 4096, op); + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 512, 4096, op, + kQwenHiddenSize, false, true); }); } else { for (const int64_t outputs : { 128, 4096 }) { @@ -8284,8 +8286,8 @@ static void register_dense_matmul_cases(Suite & suite) { [op, outputs] { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, - outputs, op); + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + outputs, op, kQwenHiddenSize, false, true); }); } } @@ -8295,8 +8297,9 @@ static void register_dense_matmul_cases(Suite & suite) { suite.device_case( "dense_matmul_swiglu.q4_k.q4_k.lowtoken.tokens" + std::to_string(tokens) + ".outputs4160", [tokens] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", - tokens, 4160); + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 4160, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, + true); }); } for (const int64_t tokens : { 1, 2, 5 }) { @@ -8304,50 +8307,34 @@ static void register_dense_matmul_cases(Suite & suite) { "dense_matmul_swiglu.q4_k.q4_k.swiglu.tokens" + std::to_string(tokens) + ".outputs16384.input4096", [tokens] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", - tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 16384, GGML_GLU_OP_SWIGLU, 4096, false, true); }); } suite.device_case("dense_matmul_swiglu.q4_k.q4_k.reglu.tokens3.outputs16384.input4352", [] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 3, 16384, - GGML_GLU_OP_REGLU, 4352); + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 3, + 16384, GGML_GLU_OP_REGLU, 4352, false, true); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.lowtoken.tokens4.outputs128", [] { run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 4, 128); + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128, + true); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs128", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 128); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + 128, true); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs4096", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 4096); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + 4096, true); }); suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens1024.outputs2048", [] { run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 1024, 2048); - }); - suite.device_case("imported_packed_gated_q4.lowtoken", [] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - 4, 128, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, true); - }); - suite.device_case("imported_packed_gated_q4.prefill", [] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", - 256, 128, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, true); - }); - suite.device_case("imported_packed_binary_q4.lowtoken", [] { - run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - 4, 128, true); - }); - suite.device_case("imported_packed_binary_q4.prefill", [] { - run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", - 256, 128, true); + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, + 2048, true); }); suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, From acecc232de79c7ef17a38879c8be32e3b85346d4 Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Tue, 3 Nov 2026 03:36:03 +0000 Subject: [PATCH 15/18] ggml-hrx: keep retained source generations fixed Remove host-write generation updates and shared atomic version storage. Preserve each allocation stamp across imported aliases and retain byte transfer and immutable weight reuse coverage. Assisted-by: Codex --- ggml/src/ggml-hrx/backend-buffer-binding.cpp | 4 +-- ggml/src/ggml-hrx/backend-context.h | 14 +--------- .../dispatch/command-program-bindings.cpp | 3 +- .../dispatch/command-program-bindings.h | 2 -- ggml/src/ggml-hrx/ggml-hrx.cpp | 16 +++-------- .../ggml-hrx/runtime/host-buffer-registry.cpp | 12 ++++---- .../ggml-hrx/runtime/host-buffer-registry.h | 12 ++++---- tests/hrx-buffer-test.cpp | 28 +++++-------------- tests/hrx-ops-test.cpp | 5 ---- 9 files changed, 27 insertions(+), 69 deletions(-) diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp index 10d03c2e04d5..3f3b7d271a77 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.cpp +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -95,7 +95,7 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx binding.offset = shared_host_region ? context->physical_offset + offset : offset; binding.length = ggml_nbytes(tensor); binding.identity = context->identity; - binding.generation = ggml_backend_hrx_buffer_generation(context); + binding.generation = context->generation; binding.capacity = shared_host_region ? context->physical_capacity : buffer->size; binding.weight = buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS; const bool retained_host_source = is_imported || is_host_buffer; @@ -117,7 +117,7 @@ bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backe return false; } region.identity = context->identity; - region.generation = ggml_backend_hrx_buffer_generation(context); + region.generation = context->generation; region.offset = context->retained_base != nullptr ? context->physical_offset + offset : offset; } else { ggml::hrx::ValueBufferBinding binding; diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index fdc8fbaf89e2..43c27ba2bbe8 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -32,26 +32,14 @@ struct ggml_backend_hrx_buffer_context { hrx_buffer_t buffer; uint8_t * base; uint64_t identity; - std::atomic generation; + uint64_t generation; bool direct_host_binding; bool imported; size_t physical_offset = 0; size_t physical_capacity = 0; uint8_t * retained_base = nullptr; - std::shared_ptr> shared_generation; }; -inline uint64_t ggml_backend_hrx_buffer_generation(const ggml_backend_hrx_buffer_context * context) { - return context->shared_generation ? context->shared_generation->load() : context->generation.load(); -} - -inline void ggml_backend_hrx_set_buffer_generation(ggml_backend_hrx_buffer_context * context, uint64_t generation) { - context->generation = generation; - if (context->shared_generation) { - *context->shared_generation = generation; - } -} - struct ggml_backend_hrx_device_context { hrx_device_t device = nullptr; std::string name; diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index a5c01da21925..d499176d14e9 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -80,7 +80,8 @@ CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBin mix_hash(hash, static_cast(static_cast(binding.value.value))); mix_hash(hash, binding.host_data != nullptr ? 1 : 0); mix_hash(hash, binding.identity); - mix_hash(hash, binding.tracks_content_generation() ? binding.generation : 0); + const bool retained_weight = binding.weight && binding.source_host_data != nullptr; + mix_hash(hash, retained_weight ? binding.generation : 0); mix_hash(hash, static_cast(binding.capacity)); mix_hash(hash, static_cast(binding.offset)); mix_hash(hash, static_cast(binding.length)); diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index ddf650ecec99..4db8e9cdb5eb 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -36,8 +36,6 @@ struct CommandProgramBinding { size_t source_length = 0; bool requires_materialization() const { return host_data != nullptr; } - // Mutable host staging refreshes at launch; retained weight content is prepared ahead of execution. - bool tracks_content_generation() const { return weight && source_host_data != nullptr; } }; inline CommandProgramBinding make_command_program_binding( diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 481af2c8f678..e2dd55a4c928 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -205,7 +205,6 @@ static void buffer_memset(ggml_backend_buffer_t buffer, GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { std::memset(context->base + destination_offset, value, size); - ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } FillBufferArgs args{ context->buffer, destination_offset, size, value }; @@ -227,7 +226,6 @@ static void buffer_set(ggml_backend_buffer_t buffer, GGML_ASSERT(destination_offset <= buffer->size && size <= buffer->size - destination_offset); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { std::memmove(context->base + destination_offset, data, size); - ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } if (!HRX_CHECK(hrx_synchronous_h2d(context->device->device, data, context->buffer, destination_offset, size))) { @@ -294,10 +292,6 @@ static bool buffer_copy(ggml_backend_buffer_t buffer, const ggml_tensor * source return false; } } - const bool destination_host = destination_context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE); - if (destination_host) { - ggml_backend_hrx_set_buffer_generation(destination_context, g_allocation_generation.fetch_add(1)); - } return true; } @@ -308,7 +302,6 @@ static void buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { auto * context = buffer_context(buffer); if (context->base != reinterpret_cast(GGML_HRX_FAKE_PTR_BASE)) { std::memset(context->base, value, buffer->size); - ggml_backend_hrx_set_buffer_generation(context, g_allocation_generation.fetch_add(1)); return; } FillBufferArgs args{ context->buffer, 0, buffer->size, value }; @@ -364,7 +357,7 @@ static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_ const uint64_t generation = g_allocation_generation.fetch_add(1); auto * context = new (std::nothrow) ggml_backend_hrx_buffer_context{ type_context->device, allocation, base, generation, generation, direct_host_binding, false, - 0, 0, nullptr, {}, + 0, 0, nullptr, }; if (context == nullptr) { if (allocation != nullptr) { @@ -375,9 +368,8 @@ static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_ if (host_visible && allocation != nullptr) { context->retained_base = base; context->physical_capacity = size; - context->shared_generation = std::make_shared>(generation); type_context->device->host_buffers.add(allocation, base, size, direct_host_binding, - context->identity, context->shared_generation); + context->identity, context->generation); } return ggml_backend_buffer_init(buft, buffer_i, context, size); } @@ -693,7 +685,7 @@ static ggml_backend_buffer_t device_buffer_from_host_ptr(ggml_backend_dev_t devi const uint64_t generation = g_allocation_generation.fetch_add(1); auto * context = new (std::nothrow) ggml_backend_hrx_buffer_context{ device_context(device), nullptr, static_cast(ptr), generation, generation, false, true, - 0, 0, nullptr, {}, + 0, 0, nullptr, }; if (context == nullptr) { return nullptr; @@ -704,7 +696,7 @@ static ggml_backend_buffer_t device_buffer_from_host_ptr(ggml_backend_dev_t devi const bool known_direct_host = known_host_region && coherent_host_region; if (known_host_region) { context->identity = known_host.identity(); - context->shared_generation = known_host.generation(); + context->generation = known_host.generation(); context->retained_base = static_cast(known_host.base()); context->physical_offset = known_host.offset(); context->physical_capacity = known_host.capacity(); diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp index fe9b04153f20..0e1b009bdf90 100644 --- a/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp @@ -9,9 +9,9 @@ namespace ggml::hrx { HostBufferRef::HostBufferRef(hrx_buffer_t buffer, size_t offset, size_t capacity, bool direct_bindable, uintptr_t base, uint64_t identity, - std::shared_ptr> generation) : + uint64_t generation) : buffer_(buffer), offset_(offset), capacity_(capacity), direct_bindable_(direct_bindable), - base_(base), identity_(identity), generation_(std::move(generation)) {} + base_(base), identity_(identity), generation_(generation) {} HostBufferRef::~HostBufferRef() { if (buffer_ != nullptr) { @@ -24,7 +24,7 @@ HostBufferRef::HostBufferRef(HostBufferRef && other) noexcept : offset_(other.offset_), capacity_(other.capacity_), direct_bindable_(other.direct_bindable_), base_(other.base_), identity_(other.identity_), - generation_(std::move(other.generation_)) { + generation_(other.generation_) { other.offset_ = 0; other.capacity_ = 0; other.direct_bindable_ = false; @@ -41,7 +41,7 @@ HostBufferRef & HostBufferRef::operator=(HostBufferRef && other) noexcept { direct_bindable_ = other.direct_bindable_; base_ = other.base_; identity_ = other.identity_; - generation_ = std::move(other.generation_); + generation_ = other.generation_; other.offset_ = 0; other.capacity_ = 0; other.direct_bindable_ = false; @@ -50,13 +50,13 @@ HostBufferRef & HostBufferRef::operator=(HostBufferRef && other) noexcept { } void HostBufferRegistry::add(hrx_buffer_t buffer, void * base, size_t size, bool direct_bindable, - uint64_t identity, std::shared_ptr> generation) { + uint64_t identity, uint64_t generation) { if (buffer == nullptr || base == nullptr || size == 0) { return; } std::lock_guard lock(mutex_); entries_.push_back({ buffer, reinterpret_cast(base), size, direct_bindable, - identity, std::move(generation) }); + identity, generation }); } void HostBufferRegistry::remove(hrx_buffer_t buffer) { diff --git a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h index 40971a87e020..a1342fe974dc 100644 --- a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h @@ -1,9 +1,7 @@ #pragma once -#include #include #include -#include #include #include @@ -31,7 +29,7 @@ class HostBufferRef { bool direct_bindable() const { return direct_bindable_; } void * base() const { return reinterpret_cast(base_); } uint64_t identity() const { return identity_; } - std::shared_ptr> generation() const { return generation_; } + uint64_t generation() const { return generation_; } private: hrx_buffer_t buffer_ = nullptr; @@ -40,10 +38,10 @@ class HostBufferRef { bool direct_bindable_ = false; uintptr_t base_ = 0; uint64_t identity_ = 0; - std::shared_ptr> generation_; + uint64_t generation_ = 0; HostBufferRef(hrx_buffer_t buffer, size_t offset, size_t capacity, bool direct_bindable, - uintptr_t base, uint64_t identity, std::shared_ptr> generation); + uintptr_t base, uint64_t identity, uint64_t generation); friend class HostBufferRegistry; }; @@ -51,7 +49,7 @@ class HostBufferRef { class HostBufferRegistry { public: void add(hrx_buffer_t buffer, void * base, size_t size, bool direct_bindable, - uint64_t identity, std::shared_ptr> generation); + uint64_t identity, uint64_t generation); void remove(hrx_buffer_t buffer); HostBufferRef find(const void * data, size_t size) const; @@ -63,7 +61,7 @@ class HostBufferRegistry { size_t size = 0; bool direct_bindable = false; uint64_t identity = 0; - std::shared_ptr> generation; + uint64_t generation = 0; }; mutable std::mutex mutex_; diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index 1c152f2627bc..b219f3e542bd 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -172,21 +172,15 @@ static void run_host_buffer_checks(ggml_backend_t backend) { ggml_backend_hrx_imported_region retained_region = {}; REQUIRE(ggml_backend_hrx_get_retained_region(host_tensor, retained_region)); REQUIRE(!ggml_backend_hrx_get_imported_region(host_tensor, retained_region)); - const uint64_t before_host_set = retained_region.generation; const uint32_t host_replacement = 0xabcdef01; ggml_backend_tensor_set(host_tensor, &host_replacement, 0, sizeof(host_replacement)); - REQUIRE(ggml_backend_hrx_get_retained_region(host_tensor, retained_region)); - REQUIRE(retained_region.generation != before_host_set); + REQUIRE(*static_cast(host_tensor->data) == host_replacement); REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - const uint64_t before_direct_set = retained_region.generation; ggml_backend_tensor_set_async(backend, direct_tensor, &host_replacement, 0, sizeof(host_replacement)); ggml_backend_synchronize(backend); - REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - REQUIRE(retained_region.generation != before_direct_set); - const uint64_t before_host_clear = retained_region.generation; + REQUIRE(*static_cast(direct_tensor->data) == host_replacement); ggml_backend_buffer_clear(direct_buffer, 0); - REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - REQUIRE(retained_region.generation != before_host_clear); + REQUIRE(*static_cast(direct_tensor->data) == 0); ggml_tensor * tensor = ggml_new_tensor_1d(ggml, GGML_TYPE_I32, 64); ggml_backend_buffer_t local = ggml_backend_alloc_buffer(backend, 4096); @@ -212,16 +206,11 @@ static void run_host_buffer_checks(ggml_backend_t backend) { REQUIRE(context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed) == upload_fallbacks); REQUIRE(context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed) == download_fallbacks); - REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - const uint64_t before_host_copy = retained_region.generation; ggml_backend_tensor_copy_async(backend, backend, tensor, direct_tensor); ggml_backend_synchronize(backend); - REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - REQUIRE(retained_region.generation != before_host_copy); - const uint64_t before_host_memset = retained_region.generation; + REQUIRE(std::memcmp(direct_tensor->data, host_words, 64 * sizeof(uint32_t)) == 0); ggml_backend_tensor_memset(direct_tensor, 0x5a, 0, sizeof(uint32_t)); - REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); - REQUIRE(retained_region.generation != before_host_memset); + REQUIRE(*static_cast(direct_tensor->data) == 0x5a5a5a5a); ggml_backend_buffer_free(local); ggml_backend_buffer_free(direct_buffer); @@ -285,12 +274,9 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { std::array readback = {}; ggml_backend_tensor_get(second, readback.data(), 0, sizeof(readback)); REQUIRE(readback[0] == source[32]); - const uint64_t old_generation = second_region.generation; const uint32_t replacement = 999; ggml_backend_tensor_set(second, &replacement, 0, sizeof(replacement)); REQUIRE(source[32] == replacement); - REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); - REQUIRE(second_region.generation != old_generation); const uint32_t async_replacement = 1234; ggml_backend_tensor_set_async(backend, second, &async_replacement, 0, sizeof(async_replacement)); @@ -298,8 +284,6 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { ggml_backend_tensor_get_async(backend, second, readback.data(), 0, sizeof(readback)); ggml_backend_synchronize(backend); REQUIRE(readback[0] == async_replacement); - REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); - REQUIRE(second_region.generation != old_generation); ggml_backend_buffer_t owned = ggml_backend_alloc_buffer(backend, sizeof(readback)); REQUIRE(owned != nullptr); @@ -332,6 +316,8 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { REQUIRE(imported_context->direct_host_binding); REQUIRE(imported_context->buffer == coherent_context->buffer); REQUIRE(imported_context->physical_offset == 256); + REQUIRE(imported_context->identity == coherent_context->identity); + REQUIRE(imported_context->generation == coherent_context->generation); ggml_init_params subparams = params; subparams.mem_size = 128 * 1024; ggml_context * subctx = ggml_init(subparams); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 93470c8ecc5c..cc2242690703 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3664,9 +3664,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, ggml_backend_hrx_cache_stats stats = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); REQUIRE(stats.weight_live_allocations == 0); - ggml_backend_hrx_imported_region input_region = {}; - REQUIRE(ggml_backend_hrx_get_retained_region(hrx_decode_input, input_region)); - const uint64_t first_input_generation = input_region.generation; compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); REQUIRE(stats.weight_materializations == 1); @@ -3677,8 +3674,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, const std::vector changed_input = make_pattern_f32(kInputSize, 41, 0.04f); set_tensor_pair_bytes(cpu_backend, cpu_decode_input, hrx_backend, hrx_decode_input, changed_input.data(), changed_input.size() * sizeof(float)); - REQUIRE(ggml_backend_hrx_get_retained_region(hrx_decode_input, input_region)); - REQUIRE(input_region.generation != first_input_generation); compute_and_compare(cpu_decode, hrx_decode, cpu_decode_output, hrx_decode_output); const std::vector changed_output = get_f32_tensor(cpu_backend, cpu_decode_output); bool output_changed = false; From 0f618f8fbcb6f219311668f977458639f8070e3e Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Tue, 3 Nov 2026 11:49:15 +0000 Subject: [PATCH 16/18] ggml-hrx: remove redundant retained-weight metadata Use existing binding bounds for retained sources, simplify the retained-region helper, and remove unused diagnostics and duplicate manual retention coverage. Assisted-by: Codex --- ggml/include/ggml-hrx.h | 5 +-- ggml/src/ggml-hrx/README.md | 2 +- ggml/src/ggml-hrx/backend-buffer-binding.cpp | 15 +-------- ggml/src/ggml-hrx/backend-buffer-binding.h | 6 ++-- .../dispatch/command-program-bindings.cpp | 2 -- .../dispatch/command-program-bindings.h | 4 --- ggml/src/ggml-hrx/ggml-hrx.cpp | 8 ++--- ggml/src/ggml-hrx/graph/value-map.cpp | 3 -- ggml/src/ggml-hrx/graph/value-map.h | 2 -- .../runtime/command-program-executor.cpp | 4 +-- .../ggml-hrx/runtime/graph-program-cache.cpp | 4 +-- ggml/src/ggml-hrx/runtime/host-memory.cpp | 1 - ggml/src/ggml-hrx/runtime/host-memory.h | 1 - tests/hrx-backend-test.cpp | 12 +++---- tests/hrx-buffer-test.cpp | 32 ++++--------------- tests/hrx-ops-test.cpp | 21 ++++++------ 16 files changed, 32 insertions(+), 90 deletions(-) diff --git a/ggml/include/ggml-hrx.h b/ggml/include/ggml-hrx.h index 2fb96acc36e4..44b29a9e79b5 100644 --- a/ggml/include/ggml-hrx.h +++ b/ggml/include/ggml-hrx.h @@ -14,12 +14,9 @@ struct ggml_backend_hrx_cache_stats { uint64_t weight_materializations; uint64_t weight_reuses; uint64_t weight_uploaded_bytes; - // Current weight-cache residency; peak bytes and materializations are cumulative. + // Current weight-cache residency. uint64_t weight_live_allocations; uint64_t weight_live_bytes; - uint64_t weight_peak_bytes; - uint64_t host_uploads; - uint64_t host_upload_bytes; uint64_t weight_duplicate_warnings; }; diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index 45b13ac6bc6b..e3152cfaf117 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -4,7 +4,7 @@ HRX materializes imported host weights in the layout requested by the selected kernel and caches each result for the backend's lifetime. If one weight is materialized in two layouts, HRX warns about the duplicate allocations and their memory cost. -To log residency totals, set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v`. HRX prints a `ggml_hrx: weight_stats` line when the backend is released, with materialization and reuse counts, cached allocations and bytes, peak cached bytes, and upload totals. +To log residency totals, set `GGML_HRX_LOG_WEIGHT_STATS=1` and run `llama-bench -v`. HRX prints a `ggml_hrx: weight_stats` line when the backend is released, with materialization and reuse counts, cached allocations and bytes, and upload totals. ## Compilation diagnostics diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp index 3f3b7d271a77..5312198b296b 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.cpp +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -77,8 +77,6 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx // Standard GGUF weights are immutable, so their native copy can use the shared weight cache. // Layout eligibility remains a separate import-only decision. binding.source_host_data = base; - binding.source_offset = host_offset; - binding.source_capacity = capacity; return true; } ggml_backend_buffer_t buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; @@ -101,12 +99,10 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx const bool retained_host_source = is_imported || is_host_buffer; binding.source_host_data = retained_host_source ? (shared_host_region ? context->retained_base : context->base) : nullptr; - binding.source_offset = binding.offset; - binding.source_capacity = binding.capacity; return true; } -bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region) { +bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_retained_region & region) { ggml_backend_hrx_buffer_context * context = nullptr; size_t offset = 0; const bool hrx_bound = ggml_backend_hrx_tensor_binding(tensor, &context, &offset); @@ -117,7 +113,6 @@ bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backe return false; } region.identity = context->identity; - region.generation = context->generation; region.offset = context->retained_base != nullptr ? context->physical_offset + offset : offset; } else { ggml::hrx::ValueBufferBinding binding; @@ -127,20 +122,12 @@ bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backe return false; } region.identity = binding.identity; - region.generation = binding.generation; region.offset = binding.offset; } region.length = ggml_nbytes(tensor); return true; } -bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region) { - ggml_backend_hrx_buffer_context * context = nullptr; - size_t offset = 0; - const bool bound = ggml_backend_hrx_tensor_binding(tensor, &context, &offset); - return bound && context->imported && ggml_backend_hrx_get_retained_region(tensor, region); -} - bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor) { if (tensor == nullptr) { return false; diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.h b/ggml/src/ggml-hrx/backend-buffer-binding.h index bd2935a36e66..e18c1260514b 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.h +++ b/ggml/src/ggml-hrx/backend-buffer-binding.h @@ -10,9 +10,8 @@ struct ggml_tensor; // The buffer identity and byte range name the source independently of a resident HRX allocation. // GGML's tensor metadata supplies the logical shape, including views into a shared import. -struct ggml_backend_hrx_imported_region { +struct ggml_backend_hrx_retained_region { uint64_t identity; - uint64_t generation; size_t offset; size_t length; }; @@ -25,8 +24,7 @@ bool ggml_backend_hrx_tensor_binding(const ggml_tensor * tensor, ggml_backend_hrx_buffer_context ** out_context, size_t * out_offset); bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx::ValueBufferBinding & binding); -bool ggml_backend_hrx_get_imported_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region); // Only borrowed imports outside HRX-owned host storage may use alternative weight layouts. bool ggml_backend_hrx_weight_has_repackable_source(const ggml_tensor * tensor); bool ggml_backend_hrx_weight_layout_eligible(const ggml_tensor * tensor); -bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_imported_region & region); +bool ggml_backend_hrx_get_retained_region(const ggml_tensor * tensor, ggml_backend_hrx_retained_region & region); diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index d499176d14e9..884175e7811d 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -88,8 +88,6 @@ CommandProgramBindingsHash command_program_bindings_hash(const CommandProgramBin mix_hash(hash, binding.weight ? 1 : 0); mix_hash(hash, binding.empty_value ? 1 : 0); mix_hash(hash, binding.source_host_data != nullptr ? 1 : 0); - mix_hash(hash, static_cast(binding.source_offset)); - mix_hash(hash, static_cast(binding.source_capacity)); } mix_hash(hash, static_cast(bindings.bindings().size())); return { hash }; diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index 4db8e9cdb5eb..20c88deea1e8 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -26,8 +26,6 @@ struct CommandProgramBinding { bool empty_value = false; // Retained host bytes are separate from the address used to bind a command. const void * source_host_data = nullptr; - size_t source_offset = 0; - size_t source_capacity = 0; // Materialized bindings are selected by the requesting command, not just by ValueId. std::string layout = kNativeWeightLayout; ggml_type layout_source_type = GGML_TYPE_COUNT; @@ -44,8 +42,6 @@ inline CommandProgramBinding make_command_program_binding( buffer.generation, buffer.capacity, buffer.host_data, buffer.weight, empty_value }; binding.source_host_data = buffer.source_host_data; - binding.source_offset = buffer.source_offset; - binding.source_capacity = buffer.source_capacity; return binding; } diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index e2dd55a4c928..6be3802203f0 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -400,10 +400,10 @@ static void backend_free(ggml_backend_t backend) { const ggml::hrx::HostWeightCacheStats weights = context->host_weights.stats(); const ggml::hrx::HostTransferStats transfers = context->host_transfers.stats(); GGML_LOG_INFO("ggml_hrx: weight_stats materializations=%" PRIu64 " reused=%" PRIu64 - " uploaded_bytes=%zu live_allocations=%zu live_bytes=%zu peak_bytes=%zu" + " uploaded_bytes=%zu live_allocations=%zu live_bytes=%zu" " host_uploads=%" PRIu64 " host_upload_bytes=%zu\n", weights.misses, weights.hits, weights.materialized_bytes, weights.allocation_count, - weights.resident_bytes, weights.peak_resident_bytes, transfers.uploads, transfers.upload_bytes); + weights.resident_bytes, transfers.uploads, transfers.upload_bytes); } context->prepared_programs.clear(); context->graph_programs.clear(); @@ -1337,7 +1337,6 @@ bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_c const ggml::hrx::GraphProgramCacheStats graph_stats = context->graph_programs.stats(); const ggml::hrx::PreparedCommandProgramCacheStats prepared_stats = context->prepared_programs.stats(); const ggml::hrx::HostWeightCacheStats weight_stats = context->host_weights.stats(); - const ggml::hrx::HostTransferStats transfer_stats = context->host_transfers.stats(); stats->graph_program_builds = graph_stats.builds; stats->graph_program_hits = graph_stats.hits; stats->prepared_program_builds = prepared_stats.builds; @@ -1347,9 +1346,6 @@ bool ggml_backend_hrx_get_cache_stats(ggml_backend_t backend, ggml_backend_hrx_c stats->weight_uploaded_bytes = weight_stats.materialized_bytes; stats->weight_live_allocations = weight_stats.allocation_count; stats->weight_live_bytes = weight_stats.resident_bytes; - stats->weight_peak_bytes = weight_stats.peak_resident_bytes; - stats->host_uploads = transfer_stats.uploads; - stats->host_upload_bytes = transfer_stats.upload_bytes; stats->weight_duplicate_warnings = weight_stats.duplicate_warnings; return true; } diff --git a/ggml/src/ggml-hrx/graph/value-map.cpp b/ggml/src/ggml-hrx/graph/value-map.cpp index ec39f0eebcfd..dfeea7c027ab 100644 --- a/ggml/src/ggml-hrx/graph/value-map.cpp +++ b/ggml/src/ggml-hrx/graph/value-map.cpp @@ -184,9 +184,6 @@ std::optional ValueMap::resolve_buffer_binding(ValueId id) c return std::nullopt; } binding.offset += value->storage_offset; - if (binding.source_host_data != nullptr) { - binding.source_offset += value->storage_offset; - } binding.length = value->byte_count; return binding; } diff --git a/ggml/src/ggml-hrx/graph/value-map.h b/ggml/src/ggml-hrx/graph/value-map.h index 46af13835a13..e41ec46719ab 100644 --- a/ggml/src/ggml-hrx/graph/value-map.h +++ b/ggml/src/ggml-hrx/graph/value-map.h @@ -65,8 +65,6 @@ struct ValueBufferBinding { bool weight = false; // Retained host bytes can back a cached native weight even when the binding is directly addressable. const void * source_host_data = nullptr; - size_t source_offset = 0; - size_t source_capacity = 0; bool requires_materialization() const { return host_data != nullptr; } }; diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 9348ad3a9bc0..7d51b7a9f9ae 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -824,8 +824,8 @@ static CommandProgramWeightDemands collect_weight_demands( demand.source.host_data = binding->source_host_data; demand.source.identity = binding->identity; demand.source.generation = binding->generation; - demand.source.capacity = binding->source_capacity; - demand.source.offset = binding->source_offset; + demand.source.capacity = binding->capacity; + demand.source.offset = binding->offset; demand.source.length = transformed ? request.source_length : binding->length; demand.source.materialized_length = transformed ? request.length : binding->length; demand.source.layout = request.layout; diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index af900bebc0e8..d62b2be3ccfb 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -491,7 +491,7 @@ void GraphProgram::bind_prepared_graph(const ggml_cgraph & graph, const GraphPro if (retained_source) { // Mutable inputs may rebind under one graph UID; immutable weight bindings are trusted. mutable_source_watches_.push_back({ match.external_bindings[i].tensor, binding.identity, - binding.source_offset, binding.length }); + binding.offset, binding.length }); } } // The full binding lookup established this graph instance's physical resources. @@ -512,7 +512,7 @@ bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { bool GraphProgram::mutable_sources_current() const { std::lock_guard lock(prepared_mutex_); for (const MutableSourceWatch & watch : mutable_source_watches_) { - ggml_backend_hrx_imported_region region = {}; + ggml_backend_hrx_retained_region region = {}; if (!ggml_backend_hrx_get_retained_region(watch.tensor, region)) { return false; } diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 73431ce93380..5a6375afca4a 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -1519,7 +1519,6 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device entries_.emplace(key, entry); ++stats_.allocation_count; stats_.resident_bytes += materialized_length; - stats_.peak_resident_bytes = std::max(stats_.peak_resident_bytes, stats_.resident_bytes); report_new_representation(key, *entry); ++stats_.misses; stats_.materialized_bytes += materialized_length; diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 87841e4baf23..23466ce7e7f4 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -72,7 +72,6 @@ struct HostWeightCacheStats { size_t materialized_bytes = 0; size_t allocation_count = 0; size_t resident_bytes = 0; - size_t peak_resident_bytes = 0; uint64_t duplicate_warnings = 0; }; diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 38beffc5e300..2355bf327a18 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1154,8 +1154,6 @@ static void run_graph_import_checks() { host_value_binding.capacity = host_storage.size(); host_value_binding.weight = true; host_value_binding.source_host_data = host_storage.data(); - host_value_binding.source_offset = 16; - host_value_binding.source_capacity = host_storage.size(); ggml::hrx::ValueMap host_values = imported.graph.values(); REQUIRE(host_values.bind_buffer(a_value->id, host_value_binding)); REQUIRE(host_values.bind_buffer(out_value->id, { dummy_hrx_buffer(0x5000), 0, out_value->byte_count })); @@ -12316,7 +12314,7 @@ static void run_command_weight_demand_checks() { binding.generation = 20; binding.weight = true; binding.source_host_data = source.data(); - binding.source_capacity = source.size(); + binding.capacity = source.size(); CommandBinding native = { "first-native", ValueId(1) }; CommandBinding packed = native; @@ -12385,13 +12383,13 @@ static void run_command_weight_demand_checks() { // Collect demand against the current source metadata and allocation eligibility. for (auto & next : current) { next.buffer = nullptr; - next.host_data = source.data() + 32; + next.host_data = source.data() + 16; next.identity = 30; next.generation = 40; next.length = 128; next.source_host_data = source.data() + 16; - next.source_offset = 32; - next.source_capacity = source.size() - 16; + next.offset = 32; + next.capacity = source.size() - 16; } const auto rebound = collect_command_program_weight_demands( commands, CommandProgramBindings::from_bindings(current)); @@ -12450,7 +12448,7 @@ static void run_command_weight_demand_validation_checks() { binding.length = source.size(); binding.weight = true; binding.source_host_data = source.data(); - binding.source_capacity = source.size(); + binding.capacity = source.size(); CommandBinding valid = { "valid", ValueId(1) }; valid.layout = "test-packed"; valid.source_type = GGML_TYPE_F32; diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index b219f3e542bd..f53d6fd41097 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -8,7 +8,6 @@ #include "ggml.h" #include "hrx-interop-utils.h" #include "runtime/host-memory.h" -#include "runtime/command-program-executor.h" #include "testing_suite.h" #include @@ -169,9 +168,8 @@ static void run_host_buffer_checks(ggml_backend_t backend) { REQUIRE(direct_binding.host_data == nullptr); REQUIRE(direct_binding.offset == 0); REQUIRE(direct_binding.length == ggml_nbytes(direct_tensor)); - ggml_backend_hrx_imported_region retained_region = {}; + ggml_backend_hrx_retained_region retained_region = {}; REQUIRE(ggml_backend_hrx_get_retained_region(host_tensor, retained_region)); - REQUIRE(!ggml_backend_hrx_get_imported_region(host_tensor, retained_region)); const uint32_t host_replacement = 0xabcdef01; ggml_backend_tensor_set(host_tensor, &host_replacement, 0, sizeof(host_replacement)); REQUIRE(*static_cast(host_tensor->data) == host_replacement); @@ -248,18 +246,18 @@ static void run_imported_buffer_checks(ggml_backend_t backend) { REQUIRE(ggml_backend_tensor_alloc(imported, second, source.data() + 32) == GGML_STATUS_SUCCESS); ggml_backend_buffer_set_usage(imported, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_hrx_imported_region first_region = {}; - ggml_backend_hrx_imported_region second_region = {}; - REQUIRE(ggml_backend_hrx_get_imported_region(first, first_region)); - REQUIRE(ggml_backend_hrx_get_imported_region(second, second_region)); + ggml_backend_hrx_retained_region first_region = {}; + ggml_backend_hrx_retained_region second_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(first, first_region)); + REQUIRE(ggml_backend_hrx_get_retained_region(second, second_region)); REQUIRE(first_region.identity == second_region.identity); REQUIRE(first_region.offset == 0); REQUIRE(second_region.offset == 32 * sizeof(uint32_t)); ggml_tensor * view = ggml_view_1d(ggml, second, 8, 4 * sizeof(uint32_t)); REQUIRE(ggml_backend_view_init(view) == GGML_STATUS_SUCCESS); - ggml_backend_hrx_imported_region view_region = {}; - REQUIRE(ggml_backend_hrx_get_imported_region(view, view_region)); + ggml_backend_hrx_retained_region view_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(view, view_region)); REQUIRE(view_region.identity == second_region.identity); REQUIRE(view_region.offset == second_region.offset + 4 * sizeof(uint32_t)); REQUIRE(view_region.length == 8 * sizeof(uint32_t)); @@ -566,26 +564,12 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(transfer_stats.uploads == 3); REQUIRE(transfer_stats.upload_bytes == 96); - ggml::hrx::RecordedCommandGraph recorded; - hrx_buffer_retain(first.lease.buffer()); - recorded.retained_buffers.push_back(first.lease.buffer()); weights.clear(); weight_stats = weights.stats(); REQUIRE(weight_stats.hits == 1); REQUIRE(weight_stats.misses == 3); REQUIRE(weight_stats.allocation_count == 0); REQUIRE(weight_stats.resident_bytes == 0); - REQUIRE(weight_stats.peak_resident_bytes == 96); - first.lease = {}; - second.lease = {}; - slice.lease = {}; - next_generation.lease = {}; - // Recording retention keeps the buffer valid independently of cache residency and weight leases. - std::array retained_bytes = {}; - require_hrx_status(hrx_synchronous_d2h(context->device->device, recorded.retained_buffers.front(), 0, - retained_bytes.data(), retained_bytes.size())); - REQUIRE(retained_bytes == first_bytes); - recorded = {}; source.layout = ggml::hrx::kNativeWeightLayout; source.generation = 3; @@ -619,7 +603,6 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(weight_stats.hits == 4); REQUIRE(weight_stats.misses == 5); REQUIRE(weight_stats.allocation_count == 2); - REQUIRE(weight_stats.peak_resident_bytes == 96); } static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { @@ -660,7 +643,6 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { REQUIRE(byte == 0xa5); } REQUIRE(weights.stats().allocation_count == 2); - REQUIRE(weights.stats().peak_resident_bytes == 2 * weight_bytes); source.layout = ggml::hrx::kNativeWeightLayout; auto native = weights.acquire(context->device->device, context->stream, transfers, source); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index cc2242690703..56ecc127fd4a 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -2861,11 +2861,11 @@ static void run_imported_native_weight_case() { ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); ggml_backend_tensor_set(reshaped_input, reshaped_input_data.data(), 0, reshaped_input_data.size() * sizeof(float)); - ggml_backend_hrx_imported_region region = {}; - REQUIRE(ggml_backend_hrx_get_imported_region(weight, region)); + ggml_backend_hrx_retained_region region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(weight, region)); REQUIRE(region.offset == 0); - ggml_backend_hrx_imported_region alias_region = {}; - REQUIRE(ggml_backend_hrx_get_imported_region(weight_alias, alias_region)); + ggml_backend_hrx_retained_region alias_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(weight_alias, alias_region)); REQUIRE(alias_region.identity == region.identity); REQUIRE(alias_region.offset == region.offset); REQUIRE(alias_region.length == region.length); @@ -3137,7 +3137,6 @@ static void run_weight_persistence_sequence_case() { for (const auto & warm : { native, packed }) { REQUIRE(warm.weight_materializations == packed_first.weight_materializations); REQUIRE(warm.weight_uploaded_bytes == packed_first.weight_uploaded_bytes); - REQUIRE(warm.weight_peak_bytes == packed_first.weight_peak_bytes); REQUIRE(warm.weight_duplicate_warnings == packed_first.weight_duplicate_warnings); REQUIRE(warm.prepared_program_builds == packed_first.prepared_program_builds); REQUIRE(warm.prepared_program_hits > packed_first.prepared_program_hits); @@ -3552,13 +3551,13 @@ static void run_imported_packed_weight_case(ggml_type weight_type, ggml::hrx::ValueBufferBinding source_binding; REQUIRE(ggml_backend_hrx_resolve_value_buffer(hrx_weight, source_binding)); REQUIRE(source_binding.source_host_data != nullptr); - REQUIRE(source_binding.source_capacity >= source.size()); + REQUIRE(source_binding.capacity >= source.size()); REQUIRE((source_binding.buffer != nullptr) == coherent_subregion); REQUIRE((source_binding.host_data != nullptr) == !coherent_subregion); if (coherent_subregion) { REQUIRE(source_binding.source_host_data == coherent_owner_base); - REQUIRE(source_binding.source_offset == 256); - REQUIRE(source_binding.source_capacity == source.size() + 256); + REQUIRE(source_binding.offset == 256); + REQUIRE(source_binding.capacity == source.size() + 256); } ggml_backend_hrx_cache_stats before_stats = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &before_stats)); @@ -3611,8 +3610,8 @@ static void run_imported_packed_weight_case(ggml_type weight_type, const auto * weight_binding = bindings.find(weight_value->id); REQUIRE(weight_binding != nullptr); REQUIRE(weight_binding->source_host_data == source_binding.source_host_data); - REQUIRE(weight_binding->source_offset == source_binding.source_offset); - REQUIRE(weight_binding->source_capacity == source_binding.source_capacity); + REQUIRE(weight_binding->offset == source_binding.offset); + REQUIRE(weight_binding->capacity == source_binding.capacity); const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); const ggml::hrx::CommandProgramExecutionContext execution_context = { backend_state->device->device, backend_state->stream, backend_state->device->architecture.c_str(), @@ -3814,7 +3813,6 @@ static void run_imported_packed_weight_case(ggml_type weight_type, REQUIRE(stats.weight_uploaded_bytes == warmed_layouts.weight_uploaded_bytes); REQUIRE(stats.weight_live_allocations == warmed_layouts.weight_live_allocations); REQUIRE(stats.weight_live_bytes == warmed_layouts.weight_live_bytes); - REQUIRE(stats.weight_peak_bytes == warmed_layouts.weight_peak_bytes); REQUIRE(stats.weight_duplicate_warnings == warmed_layouts.weight_duplicate_warnings); } REQUIRE(stats.weight_duplicate_warnings == 1); @@ -4024,7 +4022,6 @@ static void run_mixed_weight_layout_program_case(ggml_type weight_type) { REQUIRE(first.duplicate_warnings == before.duplicate_warnings + 1); REQUIRE(first.allocation_count == 2); REQUIRE(first.resident_bytes == 2 * weight.size()); - REQUIRE(first.peak_resident_bytes == first.resident_bytes); ggml_backend_hrx_cache_stats public_first = {}; REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &public_first)); REQUIRE(public_first.weight_duplicate_warnings == first.duplicate_warnings); From a2b497367f779f6ed94e8053fb7b8226e4c5feef Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Tue, 3 Nov 2026 13:27:43 +0000 Subject: [PATCH 17/18] ggml-hrx: remove retained Q6 packing scratch Use ordinary temporary materialization storage for Q6. Remove scratch-specific test setup while preserving duplicate-layout warning checks. Assisted-by: Codex --- ggml/src/ggml-hrx/runtime/host-memory.cpp | 20 +-------------- ggml/src/ggml-hrx/runtime/host-memory.h | 4 --- tests/hrx-buffer-test.cpp | 30 ++++------------------- 3 files changed, 6 insertions(+), 48 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 5a6375afca4a..71e966729f11 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -23,7 +23,6 @@ namespace { static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; static constexpr size_t kLargeHostUploadBytes = 1024 * 1024; -static constexpr size_t kQ6PackScratchLimitBytes = 64 * 1024 * 1024; struct SymmetricI4Block { std::array scales = {}; @@ -1476,25 +1475,10 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device --fail_cold_after_; } - const bool q6_packed = source.layout == kQ6KPackedK256Row64ScaleRowLayout; - const bool scratch_fits = materialized_length <= kQ6PackScratchLimitBytes; - auto scratch = q6_pack_scratch_.find(materialized_length); - const bool scratch_exists = scratch != q6_pack_scratch_.end(); - const bool scratch_has_room = materialized_length <= kQ6PackScratchLimitBytes - q6_pack_scratch_bytes_; - const bool use_scratch = q6_packed && scratch_fits && (scratch_exists || scratch_has_room); - const bool create_scratch = use_scratch && !scratch_exists; - if (create_scratch) { - scratch = q6_pack_scratch_.emplace(materialized_length, std::vector{}).first; - q6_pack_scratch_bytes_ += materialized_length; - } - const void * upload_data = nullptr; size_t upload_size = 0; std::vector transformed; - // Budget 64 MiB of host scratch across first-use materializations. The Q6 packer - // writes every output byte, and the cache lock covers packing and upload. - std::vector & output = use_scratch ? scratch->second : transformed; - result.status = materialize_weight(source, upload_data, upload_size, output); + result.status = materialize_weight(source, upload_data, upload_size, transformed); if (!result.status.success()) { return result; } @@ -1572,8 +1556,6 @@ void HostWeightCache::clear() { entries_.clear(); stats_.allocation_count = 0; stats_.resident_bytes = 0; - q6_pack_scratch_.clear(); - q6_pack_scratch_bytes_ = 0; } HostStagingBuffer::~HostStagingBuffer() { diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 23466ce7e7f4..86285f545d96 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -11,7 +11,6 @@ #include #include #include -#include typedef struct hrx_buffer_s * hrx_buffer_t; typedef struct hrx_device_s * hrx_device_t; @@ -147,9 +146,6 @@ class HostWeightCache { mutable std::mutex mutex_; std::unordered_map, SourceKeyHash> entries_; - // Size-keyed scratch is overwritten on each use and is not a source-indexed representation cache. - std::unordered_map> q6_pack_scratch_; - size_t q6_pack_scratch_bytes_ = 0; HostWeightCacheStats stats_; size_t fail_cold_after_ = SIZE_MAX; }; diff --git a/tests/hrx-buffer-test.cpp b/tests/hrx-buffer-test.cpp index f53d6fd41097..83d5ac7a4bb9 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -605,7 +605,7 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(weight_stats.allocation_count == 2); } -static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { +static void run_weight_layout_warning_checks(ggml_backend_hrx_context * context) { constexpr size_t weight_bytes = 64 * 210; std::vector host(weight_bytes, 0x5a); ggml::hrx::HostTransferManager transfers; @@ -624,26 +624,6 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { auto first = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(first.valid()); - std::vector downloaded(weight_bytes); - require_hrx_status(hrx_synchronous_d2h(context->device->device, first.lease.buffer(), 0, - downloaded.data(), downloaded.size())); - for (uint8_t byte : downloaded) { - REQUIRE(byte == 0x5a); - } - - first.lease = {}; - REQUIRE(weights.stats().allocation_count == 1); - std::fill(host.begin(), host.end(), 0xa5); - source.generation = 2; - auto second = weights.acquire(context->device->device, context->stream, transfers, source); - REQUIRE(second.valid()); - require_hrx_status(hrx_synchronous_d2h(context->device->device, second.lease.buffer(), 0, - downloaded.data(), downloaded.size())); - for (uint8_t byte : downloaded) { - REQUIRE(byte == 0xa5); - } - REQUIRE(weights.stats().allocation_count == 2); - source.layout = ggml::hrx::kNativeWeightLayout; auto native = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(native.valid()); @@ -652,7 +632,7 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { REQUIRE(alias.valid()); REQUIRE(alias.lease.buffer() == native.lease.buffer()); REQUIRE(weights.stats().duplicate_warnings == 1); - source.generation = 3; + source.generation = 2; auto changed_native = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(changed_native.valid()); REQUIRE(weights.stats().duplicate_warnings == 1); @@ -660,7 +640,7 @@ static void run_q6_pack_scratch_checks(ggml_backend_hrx_context * context) { auto changed_packed = weights.acquire(context->device->device, context->stream, transfers, source); REQUIRE(changed_packed.valid()); REQUIRE(weights.stats().duplicate_warnings == 2); - REQUIRE(weights.stats().allocation_count == 5); + REQUIRE(weights.stats().allocation_count == 4); } static void run_packed_weight_identity_checks(ggml_backend_hrx_context * context) { @@ -764,9 +744,9 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_host_weight_cache_checks(context); }); }); - suite.device_case("q6_pack_scratch", [] { + suite.device_case("weight_layout_warnings", [] { run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { - run_q6_pack_scratch_checks(context); + run_weight_layout_warning_checks(context); }); }); suite.device_case("packed_weight_identity", [] { From c4b0feb02a2a01c51203cc656ccb04b6e301c2ec Mon Sep 17 00:00:00 2001 From: AaronStGeorge Date: Tue, 3 Nov 2026 13:57:47 +0000 Subject: [PATCH 18/18] hrx: clarify value-only binding lookup Assisted-by: Codex --- .../dispatch/command-program-bindings.cpp | 2 +- .../dispatch/command-program-bindings.h | 2 +- .../runtime/command-program-executor.cpp | 2 +- tests/hrx-backend-test.cpp | 18 +++++++++--------- tests/hrx-ops-test.cpp | 4 ++-- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 884175e7811d..acc7641a47a5 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -47,7 +47,7 @@ CommandProgramBindings CommandProgramBindings::from_bindings(std::vector bindings, const Status & errors = {}); - const CommandProgramBinding * find(ValueId value) const; + const CommandProgramBinding * find_by_value(ValueId value) const; const CommandProgramBinding * find(const CommandBinding & request) const; const std::vector & bindings() const { return bindings_; } diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 7d51b7a9f9ae..b58dfc98fcda 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -1115,7 +1115,7 @@ static Status rebind_prepared_host_staging(const CommandProgramExecutionContext PreparedCommandProgram & prepared) { Status status; for (HostStagingBuffer & staging : prepared.host_staging) { - const CommandProgramBinding * binding = bindings.find(ValueId(staging.value)); + const CommandProgramBinding * binding = bindings.find_by_value(ValueId(staging.value)); if (binding == nullptr || binding->host_data == nullptr || binding->length != staging.length || binding->offset > binding->capacity || binding->length > binding->capacity - binding->offset) { status.log("live host binding does not match prepared value %d", staging.value); diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 2355bf327a18..292f389b313a 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1003,17 +1003,17 @@ static void run_graph_import_checks() { ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); REQUIRE(runtime_bindings.valid()); REQUIRE(runtime_bindings.bindings().size() == 3); - const ggml::hrx::CommandProgramBinding * a_binding = runtime_bindings.find(a_value->id); + const ggml::hrx::CommandProgramBinding * a_binding = runtime_bindings.find_by_value(a_value->id); REQUIRE(a_binding != nullptr); REQUIRE(a_binding->buffer == dummy_hrx_buffer(0x1000)); REQUIRE(a_binding->offset == 0); REQUIRE(a_binding->length == a_value->byte_count); - const ggml::hrx::CommandProgramBinding * b_binding = runtime_bindings.find(b_value->id); + const ggml::hrx::CommandProgramBinding * b_binding = runtime_bindings.find_by_value(b_value->id); REQUIRE(b_binding != nullptr); REQUIRE(b_binding->buffer == dummy_hrx_buffer(0x1800)); REQUIRE(b_binding->offset == 0); REQUIRE(b_binding->length == b_value->byte_count); - REQUIRE(runtime_bindings.find(ggml::hrx::ValueId(123456)) == nullptr); + REQUIRE(runtime_bindings.find_by_value(ggml::hrx::ValueId(123456)) == nullptr); const auto runtime_hash = ggml::hrx::command_program_bindings_hash(runtime_bindings); @@ -1160,7 +1160,7 @@ static void run_graph_import_checks() { const ggml::hrx::CommandProgramBindings host_bindings = ggml::hrx::CommandProgramBindings::from_value_map(host_values); REQUIRE(host_bindings.valid()); - const ggml::hrx::CommandProgramBinding * host_binding = host_bindings.find(a_value->id); + const ggml::hrx::CommandProgramBinding * host_binding = host_bindings.find_by_value(a_value->id); REQUIRE(host_binding != nullptr); REQUIRE(host_binding->buffer == nullptr); REQUIRE(host_binding->host_data == host_storage.data()); @@ -1172,7 +1172,7 @@ static void run_graph_import_checks() { REQUIRE(host_values.bind_buffer(a_value->id, staging_value_binding)); const auto staging_bindings = ggml::hrx::CommandProgramBindings::from_value_map(host_values); REQUIRE(staging_bindings.valid()); - const auto * staging_binding = staging_bindings.find(a_value->id); + const auto * staging_binding = staging_bindings.find_by_value(a_value->id); REQUIRE(staging_binding != nullptr); REQUIRE(staging_binding->host_data == host_storage.data()); REQUIRE(staging_binding->source_host_data == nullptr); @@ -11169,7 +11169,7 @@ static void run_alias_value_import_checks() { const ggml::hrx::CommandProgramBindings bindings = ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); REQUIRE(bindings.valid()); - const ggml::hrx::CommandProgramBinding * view_binding = bindings.find(view_value->id); + const ggml::hrx::CommandProgramBinding * view_binding = bindings.find_by_value(view_value->id); REQUIRE(view_binding != nullptr); REQUIRE(view_binding->buffer == dummy_hrx_buffer(0x4000)); REQUIRE(view_binding->offset == 128 + 2 * sizeof(float)); @@ -11220,7 +11220,7 @@ static void run_alias_value_import_checks() { const ggml::hrx::CommandProgramBindings internal_bindings = ggml::hrx::CommandProgramBindings::from_value_map(internal_imported.graph.values()); REQUIRE(internal_bindings.valid()); - const ggml::hrx::CommandProgramBinding * internal_view_binding = internal_bindings.find(internal_view_value->id); + const ggml::hrx::CommandProgramBinding * internal_view_binding = internal_bindings.find_by_value(internal_view_value->id); REQUIRE(internal_view_binding != nullptr); REQUIRE(internal_view_binding->buffer == dummy_hrx_buffer(0x5000)); REQUIRE(internal_view_binding->offset == 256 + internal_source->nb[1]); @@ -11743,7 +11743,7 @@ static void run_zero_output_unused_weight_staging(bool direct_execution) { } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); REQUIRE(bindings.valid()); - REQUIRE(bindings.find(unused_value->id) != nullptr); + REQUIRE(bindings.find_by_value(unused_value->id) != nullptr); const ggml::hrx::CommandProgramExecutionContext execution_context = { state->device->device, state->stream, state->device->architecture.c_str(), &corpus, &state->kernel_executables, &state->transient_arena, &state->host_transfers, @@ -12019,7 +12019,7 @@ static void run_chained_dispatch_requires_transients() { const ggml::hrx::CommandProgramBindings bindings = ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); REQUIRE(bindings.valid()); - REQUIRE(bindings.find(sum_value->id) == nullptr); + REQUIRE(bindings.find_by_value(sum_value->id) == nullptr); const ggml::hrx::ResolvedCommandProgram resolved = ggml::hrx::resolve_command_program_bindings(commands, bindings); REQUIRE(!resolved.valid()); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 56ecc127fd4a..feed4d3ed733 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -3607,7 +3607,7 @@ static void run_imported_packed_weight_case(ggml_type weight_type, } auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); REQUIRE(bindings.valid()); - const auto * weight_binding = bindings.find(weight_value->id); + const auto * weight_binding = bindings.find_by_value(weight_value->id); REQUIRE(weight_binding != nullptr); REQUIRE(weight_binding->source_host_data == source_binding.source_host_data); REQUIRE(weight_binding->offset == source_binding.offset); @@ -4096,7 +4096,7 @@ static void run_mixed_weight_layout_program_case(ggml_type weight_type) { REQUIRE(bindings.valid()); const auto * ids_value = values.find_tensor(hrx_ids); REQUIRE(ids_value != nullptr); - const auto * ids_binding = bindings.find(ids_value->id); + const auto * ids_binding = bindings.find_by_value(ids_value->id); REQUIRE(ids_binding != nullptr && !ids_binding->weight); const auto & selected = lookup.program->commands(); bool found_packed_weight_request = false;