diff --git a/ggml/include/ggml-hrx.h b/ggml/include/ggml-hrx.h index d40c095b75b..44b29a9e79b 100644 --- a/ggml/include/ggml-hrx.h +++ b/ggml/include/ggml-hrx.h @@ -11,6 +11,13 @@ 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; + // Current weight-cache residency. + uint64_t weight_live_allocations; + uint64_t weight_live_bytes; + uint64_t weight_duplicate_warnings; }; GGML_BACKEND_API ggml_backend_t ggml_backend_hrx_init(size_t device); diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index b60ec1509a9..e3152cfaf11 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -1,5 +1,11 @@ # HRX Backend +## Weight residency diagnostics + +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, and upload totals. + ## Compilation diagnostics Set `GGML_HRX_TRACE_COMPILE=1` before starting an application using the ggml HRX backend to log compilation activity. Each record is a JSON object on stderr prefixed with `hrx compile: `. This output bypasses the application's log-level filter, so it also works with quiet benchmark runs. The setting is read once per process; unset or any value other than `1` disables it. diff --git a/ggml/src/ggml-hrx/backend-buffer-binding.cpp b/ggml/src/ggml-hrx/backend-buffer-binding.cpp index bab5f12b1af..5312198b296 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.cpp +++ b/ggml/src/ggml-hrx/backend-buffer-binding.cpp @@ -27,7 +27,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; @@ -68,19 +74,86 @@ bool ggml_backend_hrx_resolve_value_buffer(const ggml_tensor * tensor, ggml::hrx 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. + // Layout eligibility remains a separate import-only decision. + binding.source_host_data = base; 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.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; return true; } + +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); + 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.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.offset = binding.offset; + } + region.length = ggml_nbytes(tensor); + return true; +} + +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 59b925873f6..e18c1260514 100644 --- a/ggml/src/ggml-hrx/backend-buffer-binding.h +++ b/ggml/src/ggml-hrx/backend-buffer-binding.h @@ -8,6 +8,14 @@ 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_retained_region { + uint64_t identity; + size_t offset; + size_t length; +}; + 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 +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); +// 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_retained_region & region); diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index 3d510b310b5..43c27ba2bbe 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -34,6 +34,10 @@ struct ggml_backend_hrx_buffer_context { uint64_t identity; uint64_t generation; bool direct_host_binding; + bool imported; + size_t physical_offset = 0; + size_t physical_capacity = 0; + uint8_t * retained_base = nullptr; }; struct ggml_backend_hrx_device_context { diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 741a41c9c17..acc7641a47a 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 { @@ -27,9 +26,7 @@ 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 }); + bindings.push_back(make_command_program_binding(value->id, *buffer, value->byte_count == 0)); } return from_bindings(std::move(bindings), result.status); } @@ -50,7 +47,7 @@ CommandProgramBindings CommandProgramBindings::from_bindings(std::vector(static_cast(binding.value.value))); mix_hash(hash, binding.host_data != nullptr ? 1 : 0); mix_hash(hash, binding.identity); - mix_hash(hash, binding.generation); + 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)); 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(bindings.bindings().size())); 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.generation - << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length - << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0); - } - 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 d43cc5e8f3e..8cd57c904dc 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -1,6 +1,7 @@ #pragma once #include "graph/value-map.h" +#include "command-program.h" #include "status.h" #include @@ -23,13 +24,28 @@ 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; + // 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; } }; -struct CommandProgramBindingsFingerprint { - std::string 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; + return binding; +} + + struct CommandProgramBindingsHash { uint64_t value = 0; @@ -41,7 +57,8 @@ class CommandProgramBindings { static 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_; } @@ -54,6 +71,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/dispatch/command-program-resolver.cpp b/ggml/src/ggml-hrx/dispatch/command-program-resolver.cpp index 18805390c90..dfc41d6a676 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 c0016ba9f7d..cfeb31398a7 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 047caa6acec..4e785ac0fb7 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 193f3dac977..7555e200fa6 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 f491890c1bd..2fe76c61798 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 98d90b2957d..4e07e2933ad 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 0f21047a92d..6be3802203f 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) { @@ -222,7 +225,7 @@ 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); return; } if (!HRX_CHECK(hrx_synchronous_h2d(context->device->device, data, context->buffer, destination_offset, size))) { @@ -267,8 +270,29 @@ 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; + } + } + return true; } static void buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { @@ -332,7 +356,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 +366,10 @@ 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; + type_context->device->host_buffers.add(allocation, base, size, direct_host_binding, + context->identity, context->generation); } return ggml_backend_buffer_init(buft, buffer_i, context, size); } @@ -368,6 +396,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" + " host_uploads=%" PRIu64 " host_upload_bytes=%zu\n", + weights.misses, weights.hits, weights.materialized_bytes, weights.allocation_count, + weights.resident_bytes, transfers.uploads, transfers.upload_bytes); + } context->prepared_programs.clear(); context->graph_programs.clear(); context->kernel_executables.clear(); @@ -436,6 +473,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 +503,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 +524,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 +537,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); } @@ -562,7 +635,8 @@ 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 }; + props->caps.mmap_support = true; } static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parameters) { @@ -584,6 +658,7 @@ static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parame if (backend == nullptr) { delete context; hrx_stream_release(stream); + } else { } return backend; } @@ -596,6 +671,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->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 @@ -1090,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, @@ -1218,13 +1333,20 @@ 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); 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(); 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; + stats->weight_live_allocations = weight_stats.allocation_count; + stats->weight_live_bytes = weight_stats.resident_bytes; + stats->weight_duplicate_warnings = weight_stats.duplicate_warnings; return true; } diff --git a/ggml/src/ggml-hrx/graph/value-map.h b/ggml/src/ggml-hrx/graph/value-map.h index b026de65fde..e41ec46719a 100644 --- a/ggml/src/ggml-hrx/graph/value-map.h +++ b/ggml/src/ggml-hrx/graph/value-map.h @@ -63,6 +63,8 @@ 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; 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 2b03090b48f..b58dfc98fcd 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); @@ -699,18 +703,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,44 +730,116 @@ static std::unordered_map collect_graph_value_access( access.write = true; break; } - if (binding.access == ResourceAccess::Write) { + } + } + }; + append_command_list_access(commands.initialization_commands, false); + append_command_list_access(commands.commands, true); + 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 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; } - if (binding.layout == kNativeWeightLayout) { - access.native_read = true; - if (access.transformed_read) { - access.layout_conflict = true; - } + 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; } - 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; + 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; } - 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; + 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; + } } - if (access.native_read) { - access.layout_conflict = true; + + // 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->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; + 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)); } } }; - append_command_list_access(commands.initialization_commands, false); - append_command_list_access(commands.commands, true); - return access_by_value; + collect(commands.initialization_commands); + collect(commands.commands); + return result; } static CommandProgramBindings materialize_host_bindings(const CommandProgramExecutionContext & context, @@ -783,57 +850,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_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); 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; } @@ -1053,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); @@ -1779,6 +1841,11 @@ static bool execute_prepared_command_prefix_via_graph(const CommandProgramExecut } // namespace +CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, + const CommandProgramBindings & bindings) { + return collect_weight_demands(commands, bindings, collect_graph_value_access(commands)); +} + bool debug_serial_command_execution_enabled() { return environment_flag_enabled("GGML_HRX_DEBUG_SERIAL_EXECUTION") || dispatch_trace_enabled(); } @@ -2169,6 +2236,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 f2035a88698..7a35a01080a 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -46,6 +46,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: @@ -66,6 +67,22 @@ struct CommandProgramExecutionContext { CompileTraceExecution * compile_trace = nullptr; }; +struct CommandProgramWeightDemand { + CommandBinding request; + HostWeightSource source; + bool requires_materialization = false; +}; + +struct CommandProgramWeightDemands { + std::vector weights; + Status status; + + bool valid() const { return status.success(); } +}; + +CommandProgramWeightDemands collect_command_program_weight_demands(const CommandProgram & commands, + const CommandProgramBindings & bindings); + struct PreparedCommandBinding { CommandBinding binding; ResolvedBufferRef ref; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 4d1e7db8096..d192393596a 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -2,6 +2,8 @@ #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" @@ -51,22 +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)) { - 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; + 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); } @@ -95,14 +87,17 @@ 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 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); - 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; - } + if (reusable_owner) { + 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); + return result; } CommandProgramBindings bindings = bind_external_value_buffers(binding_match); @@ -123,11 +118,18 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { &context_.graph_replay_state, trace.ids.cache_id != 0 ? &trace : nullptr, }; - const PreparedCommandProgramCacheExecutionResult execution = - use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings) : - 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.success) { result.status.append(execution.status); if (result.status.success()) { @@ -136,6 +138,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, bindings, std::move(entry)); + } + result.code = GGML_STATUS_SUCCESS; trace.success = true; return result; diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.h b/ggml/src/ggml-hrx/runtime/graph-executor.h index a2759ea7187..7b66ae160cc 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 666602f8fcd..d62b2be3ccf 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -1,10 +1,12 @@ #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 #include #include #include @@ -157,21 +159,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) { @@ -407,6 +394,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; @@ -415,12 +406,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(prepared_.host_staging.size()); - for (const HostStagingBuffer & staging : prepared_.host_staging) { + std::lock_guard entry_lock(prepared_entry_->mutex); + 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)); if (slot == nullptr) { result.status.log("prepared host staging value %d has no external graph slot", staging.value); @@ -439,11 +432,11 @@ Status GraphProgram::capture_external_slots(const ggml_cgraph & graph, const Gra Status status; external_slots_.clear(); external_slot_by_value_.clear(); - fast_path_nodes_ = graph.nodes; external_slots_.reserve(match.external_bindings.size()); for (const GraphProgramExternalBinding & binding : match.external_bindings) { 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]; @@ -479,83 +472,62 @@ 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 has_prepared_; + return prepared_entry_; } -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_; +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(); + // 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({ match.external_bindings[i].tensor, binding.identity, + binding.offset, binding.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; } -PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { +bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { std::lock_guard lock(prepared_mutex_); - return prepared_stats_; + // 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 == fast_path_uid_; + const bool same_nodes = graph.nodes == fast_path_nodes_; + return allocated_graph && same_uid && same_nodes; } -PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( - const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings) { - PreparedCommandProgramCacheExecutionResult result; - 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; - } - +bool GraphProgram::mutable_sources_current() const { std::lock_guard lock(prepared_mutex_); - if (!has_prepared_) { - prepared_ = prepare_command_program(context, *commands_, bindings); - if (!prepared_.valid()) { - result.status.append(prepared_.status); - return result; - } - has_prepared_ = true; - ++prepared_stats_.builds; - } else { - ++prepared_stats_.hits; - } - - 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_); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); + for (const MutableSourceWatch & watch : mutable_source_watches_) { + ggml_backend_hrx_retained_region region = {}; + if (!ggml_backend_hrx_get_retained_region(watch.tensor, region)) { + return false; } - return result; - } - - if (environment_flag_enabled("GGML_HRX_REBUILD_GRAPH_REPLAY")) { - recorded_ = {}; - } - - const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, *commands_, bindings, prepared_, 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"); + 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 identity_changed = !same_identity; + const bool offset_changed = !same_offset; + const bool length_changed = !same_length; + const bool source_changed = identity_changed || offset_changed || length_changed; + if (source_changed) { + return false; } - return result; } - result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, prepared_); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } - return result; + return true; } GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_graph) const { @@ -650,7 +622,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; } @@ -754,7 +735,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); @@ -819,13 +800,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 4cdb3774370..54c41ca6873 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,7 @@ #include #include #include +#include #include struct ggml_cgraph; @@ -33,6 +35,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 +67,17 @@ class GraphProgram { const CommandProgram & commands() const { return *commands_; } - CommandProgram & commands() { return *commands_; } - 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; - 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; - - PreparedCommandProgramCacheStats prepared_stats() const; - - PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, - const CommandProgramBindings & bindings); + bool mutable_sources_current() 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: const GraphProgramExternalSlot * find_external_slot(ValueId value) const; @@ -89,25 +88,28 @@ class GraphProgram { uint64_t uid_ = 0; std::string target_; std::unique_ptr graph_; - std::unique_ptr commands_; + std::unique_ptr commands_; std::string command_shape_; uint64_t command_shape_hash_ = 0; std::vector external_slots_; + struct MutableSourceWatch { + const ggml_tensor * tensor; + uint64_t identity; + size_t offset; + size_t length; + }; + std::vector mutable_source_watches_; 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_; - PreparedCommandProgram prepared_; - RecordedCommandGraph recorded_; - 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/graph-replay.h b/ggml/src/ggml-hrx/runtime/graph-replay.h index 3bcd3abdcbb..5a9cd4ed414 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/host-buffer-registry.cpp b/ggml/src/ggml-hrx/runtime/host-buffer-registry.cpp index 02b591fcfc9..0e1b009bdf9 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, + uint64_t generation) : + buffer_(buffer), offset_(offset), capacity_(capacity), direct_bindable_(direct_bindable), + base_(base), identity_(identity), generation_(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_(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_ = 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, uint64_t 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, 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 30ce0eba2ec..a1342fe974d 100644 --- a/ggml/src/ggml-hrx/runtime/host-buffer-registry.h +++ b/ggml/src/ggml-hrx/runtime/host-buffer-registry.h @@ -25,19 +25,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_; } + uint64_t 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; + uint64_t generation_ = 0; + + HostBufferRef(hrx_buffer_t buffer, size_t offset, size_t capacity, bool direct_bindable, + uintptr_t base, uint64_t identity, uint64_t 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, uint64_t generation); void remove(hrx_buffer_t buffer); HostBufferRef find(const void * data, size_t size) const; @@ -47,6 +59,9 @@ class HostBufferRegistry { hrx_buffer_t buffer = nullptr; uintptr_t base = 0; size_t size = 0; + bool direct_bindable = false; + uint64_t identity = 0; + uint64_t generation = 0; }; 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 c1741fd802e..71e966729f1 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,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -22,14 +24,6 @@ namespace { static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; static constexpr size_t kLargeHostUploadBytes = 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 = {}; @@ -1258,13 +1252,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, @@ -1368,55 +1355,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_; @@ -1492,10 +1430,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; } @@ -1511,55 +1453,32 @@ 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 }; - { - std::lock_guard lock(mutex_); - const auto found = entries_.find(key); - if (found != entries_.end()) { - ++stats_.hits; - result.lease = HostWeightLease(found->second); - return result; - } + // 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()) { + ++stats_.hits; + result.lease = HostWeightLease(found->second); + return result; } - - 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; + 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_; } const void * upload_data = nullptr; size_t upload_size = 0; std::vector transformed; - result.status = materialize_weight(materialization_source, upload_data, upload_size, transformed); - if (mapped_canonical.lock.owns_lock()) { - mapped_canonical.lock.unlock(); - } + result.status = materialize_weight(source, upload_data, upload_size, transformed); if (!result.status.success()) { return result; } @@ -1581,18 +1500,12 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device 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_.allocation_count; + stats_.resident_bytes += materialized_length; + report_new_representation(key, *entry); + ++stats_.misses; + stats_.materialized_bytes += materialized_length; result.lease = HostWeightLease(std::move(entry)); return result; } @@ -1602,10 +1515,47 @@ HostWeightCacheStats HostWeightCache::stats() const { return stats_; } +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; + 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_key.length == key.length; + const bool same_region = same_identity && same_generation && same_capacity && same_offset && + same_length; + if (!same_region) { + continue; + } + 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 + key.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) { + std::lock_guard lock(mutex_); + fail_cold_after_ = successful_acquires; +} + void HostWeightCache::clear() { std::lock_guard lock(mutex_); entries_.clear(); - stats_ = {}; + stats_.allocation_count = 0; + stats_.resident_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 a125598f20b..86285f545d9 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -25,15 +25,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, @@ -49,27 +42,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; @@ -82,11 +64,14 @@ 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; + size_t materialized_bytes = 0; size_t allocation_count = 0; size_t resident_bytes = 0; + uint64_t duplicate_warnings = 0; }; class HostWeightLease { @@ -113,6 +98,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() = default; @@ -126,6 +112,8 @@ class HostWeightCache { HostTransferManager & transfers, const HostWeightSource & source); HostWeightCacheStats stats() const; + void fail_cold_acquire_after_for_test(size_t successful_acquires); + // Clears cache residency; existing leases remain valid and cumulative diagnostics are preserved. void clear(); private: @@ -154,9 +142,12 @@ 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_; HostWeightCacheStats stats_; + 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 ad81b1a23af..5ab1fcc193d 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -1,7 +1,9 @@ #include "prepared-command-program-cache.h" +#include #include #include +#include #include #include @@ -98,11 +100,12 @@ 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); + auto prepared = prepare_command_program(context, commands, bindings); if (!prepared.valid()) { result.status.append(prepared.status); return result; } + // 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, prepared); if (!result.success) { result.status.log("execute uncached HRX command program failed"); @@ -110,96 +113,74 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ return result; } - 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; - } + 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); - if (entry->has_program && entry->program.valid()) { - 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); - 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); - 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"); - } + 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->ready) { + auto prepared = prepare_command_program(context, commands, bindings); + if (!prepared.valid()) { + result.status.append(prepared.status); return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); - if (!result.success) { - result.status.log("execute cached HRX command program failed"); - } - return result; + entry->prepared = std::move(prepared); + entry->ready = true; + record_build(); + } else { + record_hit(); } - 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); - return result; - } - entry->program = std::move(prepared); - entry->has_program = true; - 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); - if (!result.success) { - result.status.log("execute prepared HRX command program failed"); + 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) { + entry->recorded = {}; } - return result; - } - - const RecordedCommandGraphExecutionResult replay = - bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->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"); + const auto replay = bind_and_launch_recorded_command_graph( + context, commands, bindings, entry->prepared, entry->recorded); + apply_graph_replay_result(result, replay); + 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); } - return result; } - result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); - 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 069ffd62448..d18ef44375e 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 { @@ -21,6 +22,14 @@ struct PreparedCommandProgramCacheStats { uint64_t hits = 0; }; +struct PreparedCommandProgramCacheEntry { + std::mutex mutex; + // Release recorded handles before the prepared bindings they reference. + PreparedCommandProgram prepared; + RecordedCommandGraph recorded; + bool ready = false; +}; + struct PreparedCommandProgramCacheExecutionResult { bool success = false; Status status; @@ -35,6 +44,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, @@ -79,18 +101,11 @@ class PreparedCommandProgramCache { uint64_t command_shape_hash, const CommandProgramBindings & bindings) const; - struct Entry { - std::mutex mutex; - PreparedCommandProgram program; - RecordedCommandGraph recorded; - bool has_program = false; - }; - 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 659cff9bb0a..292f389b313 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -1003,21 +1003,19 @@ 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 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,17 +1023,20 @@ 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); - - 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(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(); + 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 != - runtime_fingerprint.value); + REQUIRE(ggml::hrx::command_program_bindings_hash(changed_generation_bindings).value == + runtime_hash.value); ggml::hrx::ValueMap changed_capacity = imported.graph.values(); REQUIRE(changed_capacity.bind_buffer( @@ -1043,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"); @@ -1152,20 +1153,40 @@ 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(); 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 })); 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()); REQUIRE(host_binding->offset == 16); REQUIRE(host_binding->length == a_value->byte_count); REQUIRE(host_binding->weight); - REQUIRE(ggml::hrx::command_program_bindings_fingerprint(host_bindings).value != runtime_fingerprint.value); + 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_by_value(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_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; }); + 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); resolved = ggml::hrx::resolve_command_program_bindings(commands, missing_bindings); REQUIRE(!resolved.valid()); @@ -6163,6 +6184,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 +6243,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 +6325,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 +6399,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 +6758,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())); @@ -11139,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)); @@ -11190,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]); @@ -11445,16 +11475,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 +11501,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 +11667,146 @@ 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->byte_count == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external_bindings)); + REQUIRE(bindings.valid()); + 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, + &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 = 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 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(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(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); + 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); @@ -11838,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()); @@ -12122,6 +12303,232 @@ static void run_disjoint_transient_plan_packing_checks() { ggml_free(ctx); } +static void run_command_weight_demand_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.capacity = source.size(); + + 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); + 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 = collect_command_program_weight_demands(commands, 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); + // Collect demand against the current source metadata and allocation eligibility. + for (auto & next : current) { + next.buffer = nullptr; + next.host_data = source.data() + 16; + next.identity = 30; + next.generation = 40; + next.length = 128; + next.source_host_data = source.data() + 16; + next.offset = 32; + next.capacity = source.size() - 16; + } + 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) { + 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.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_COUNT); + 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 = 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); + auto uploaded = binding; + uploaded.source_host_data = nullptr; + 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 = 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 = 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_demand_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.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 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); + 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 = 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); + for (const auto & error : missing_source.status.errors()) { + REQUIRE(error.find("no retained source") != std::string::npos); + } + missing.weight = false; + const auto nonweight = collect_command_program_weight_demands( + commands, CommandProgramBindings::from_bindings({ missing })); + REQUIRE(nonweight.valid()); + REQUIRE(nonweight.weights.empty()); + const auto unbound = collect_command_program_weight_demands(commands, 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); + auto short_binding = binding; + short_binding.length = valid.source_length - 1; + 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 = collect_command_program_weight_demands( + commands, 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 +12537,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())); @@ -12534,6 +12942,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 = {}; @@ -14135,6 +14656,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_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(); }); @@ -14214,10 +14737,74 @@ 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(); }); + 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-buffer-test.cpp b/tests/hrx-buffer-test.cpp index e281a3199f0..83d5ac7a4bb 100644 --- a/tests/hrx-buffer-test.cpp +++ b/tests/hrx-buffer-test.cpp @@ -2,18 +2,24 @@ #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 "testing_suite.h" +#include #include #include #include #include #include #include +#if defined(__linux__) +#include +#endif #define REQUIRE(condition) \ do { \ @@ -162,6 +168,17 @@ 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_retained_region retained_region = {}; + REQUIRE(ggml_backend_hrx_get_retained_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); + REQUIRE(ggml_backend_hrx_get_retained_region(direct_tensor, retained_region)); + ggml_backend_tensor_set_async(backend, direct_tensor, &host_replacement, 0, sizeof(host_replacement)); + ggml_backend_synchronize(backend); + REQUIRE(*static_cast(direct_tensor->data) == host_replacement); + ggml_backend_buffer_clear(direct_buffer, 0); + 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); @@ -187,12 +204,214 @@ 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); + ggml_backend_tensor_copy_async(backend, backend, tensor, direct_tensor); + ggml_backend_synchronize(backend); + 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(*static_cast(direct_tensor->data) == 0x5a5a5a5a); + 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_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_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)); + + 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 uint32_t replacement = 999; + ggml_backend_tensor_set(second, &replacement, 0, sizeof(replacement)); + REQUIRE(source[32] == replacement); + + 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); + + 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); + 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); + 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; + REQUIRE(ggml_backend_graph_compute(backend, coherent_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + 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; @@ -347,10 +566,142 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { 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 == 0); REQUIRE(weight_stats.resident_bytes == 0); + + source.layout = ggml::hrx::kNativeWeightLayout; + source.generation = 3; + source.offset = 16; + source.source_type = GGML_TYPE_F32; + 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()); + // 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()); + 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 == 4); + REQUIRE(weight_stats.misses == 5); + REQUIRE(weight_stats.allocation_count == 2); +} + +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; + 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()); + 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 = 2; + 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 == 4); +} + +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 @@ -373,6 +724,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 +744,16 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_host_weight_cache_checks(context); }); }); + suite.device_case("weight_layout_warnings", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_weight_layout_warning_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 73d5f79efc8..feed4d3ed73 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,1569 @@ 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); + 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)); + 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_tensor_set(reshaped_input, reshaped_input_data.data(), 0, reshaped_input_data.size() * sizeof(float)); + + ggml_backend_hrx_retained_region region = {}; + REQUIRE(ggml_backend_hrx_get_retained_region(weight, region)); + REQUIRE(region.offset == 0); + 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); + 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); + } + // 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); + 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->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_persistence_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) { + 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, 1); + REQUIRE(native_first.weight_materializations == 1); + 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_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_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()); + REQUIRE(recorded_native.graph_replay_event == (replay_disabled ? ggml::hrx::HrxGraphReplayEvent::Disabled : + ggml::hrx::HrxGraphReplayEvent::Hit)); + ggml_backend_synchronize(backend); + + const auto structural_before = state->prepared_programs.stats(); + 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, 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); + REQUIRE(hrx_no_weight->n_nodes > 0); + 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); + 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); + REQUIRE(state->host_weights.stats().allocation_count == 3); + }; + auto fail_then_retry = [&](ggml_cgraph * two_weights, ggml_tensor * first, ggml_tensor * second) { + ggml_backend_synchronize(backend); + 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); + ggml_backend_synchronize(backend); + 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); + verify_two_weights(two_weights, first, second); + REQUIRE(state->host_weights.stats().misses == failed.weight_materializations + 1); + }; + 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.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); + 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_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"); + 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_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_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_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + hrx_graph->uid = 951081; + + 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 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); + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + 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)); + 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)); + 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); + REQUIRE(packed_lookup.valid()); + auto packed_bindings = bind(*packed_lookup.program, *hrx_graph); + 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); + const auto packed_source = packed_demands.weights.front().source; + 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); + 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, + }; + 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. + 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 }))); + 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 bool execution_failed = !ggml::hrx::execute_prepared_command_program( + execution_context, *prepared); + std::weak_ptr sole_resource = entry; + const bool flush_ok = !ggml::hrx::take_status(hrx_stream_flush(state->stream)); + entry.reset(); + uncached_prepared.reset(); + if (uncached) { + state->graph_programs.clear(); + state->kernel_executables.clear(); + } + const auto pending_stats = state->host_weights.stats(); + bool stream_complete = true; + const bool query_ok = !ggml::hrx::take_status(hrx_stream_query(state->stream, &stream_complete)); + const bool cache_only_owner = sole_resource.use_count() == 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(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); + 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(); + state->prepared_programs.clear(); + REQUIRE(sole_resource.expired()); + 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); + 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_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); + 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 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(), + 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.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.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)); + 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)); + external.push_back(ggml::hrx::make_command_program_binding( + slot.value, value, ggml_nbytes(slot.tensor) == 0)); + } + auto bindings = ggml::hrx::CommandProgramBindings::from_bindings(std::move(external)); + REQUIRE(bindings.valid()); + 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); + 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(), + &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); + 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)); + 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->mutable_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)); + external.push_back(ggml::hrx::make_command_program_binding( + slot.value, value, 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 = 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'; + 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); + 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::graph_replay_disabled_by_environment() ? ggml::hrx::HrxGraphReplayEvent::Disabled : + 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 = 2; + REQUIRE(stats.weight_materializations == native_materializations); + 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 == 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_duplicate_warnings == warmed_layouts.weight_duplicate_warnings); + } + REQUIRE(stats.weight_duplicate_warnings == 1); + + 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->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)); + 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->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); + 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); + + { + // 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(); + 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)); + 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()); + 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); + 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(get_f32_tensor(hrx_backend, hrx_decode_output) == expected); + 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); + 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()); + 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_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); + 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_by_value(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_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); + + // 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 = {}; + 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); @@ -4209,6 +5824,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, @@ -4216,7 +5841,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); @@ -4287,6 +5913,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; @@ -4302,7 +5931,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); @@ -4323,6 +5964,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); @@ -4503,7 +6149,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); @@ -4531,8 +6179,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); @@ -4547,20 +6199,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); @@ -4570,6 +6253,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, @@ -4637,7 +6387,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); @@ -4666,8 +6417,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); @@ -4682,14 +6437,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); @@ -4737,8 +6498,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); @@ -4789,13 +6559,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); @@ -4888,7 +6655,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 }); } @@ -6391,15 +8159,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); }); @@ -6489,7 +8257,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_q4_q8_1_x4_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 }) { @@ -6499,7 +8268,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_swiglu_q4_k_f16_wmma_prefill_wave32", 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 }) { @@ -6509,7 +8279,7 @@ static void register_dense_matmul_cases(Suite & suite) { 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, - outputs, op); + outputs, op, kQwenHiddenSize, false, true); }); } } @@ -6520,7 +8290,8 @@ static void register_dense_matmul_cases(Suite & suite) { "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", - tokens, 4160); + tokens, 4160, GGML_GLU_OP_SWIGLU, kQwenHiddenSize, false, + true); }); } for (const int64_t tokens : { 1, 2, 5 }) { @@ -6529,29 +8300,33 @@ static void register_dense_matmul_cases(Suite & suite) { [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", - tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); + 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_q4_q8_1_x4_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_q4_q8_1_x4_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_swiglu_q4_k_f16_wmma_prefill_wave32", 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_swiglu_q4_k_f16_wmma_prefill_wave32", 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_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); + 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, @@ -6662,11 +8437,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); }); } @@ -6749,6 +8524,40 @@ 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.device_case("weight_persistence_sequence", [] { run_weight_persistence_sequence_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); }); + 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("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 }) {