From de659677f031ebbbf00ce4d07525e3ca05738fc4 Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Wed, 23 Sep 2026 21:56:30 -0700 Subject: [PATCH 1/8] Factor HRX vector matmul into Loom kernels Move ordinary batch-size-1 matmul decode coverage to vector-shaped Loom kernels, including vector postops, packed Q6, selected-output Q6, and optional publish paths. Keep Qwen RMSNorm plus Q6 vocab projection on its fused Q8 activation route and leave skinny/tiled ownership with existing paths. Assisted-by: Codex --- .../common/dispatch-mul-mat.cpp | 776 +++++++-- .../kernels/loom-libs/manifest.json | 399 ++++- .../motifs/matmul/vector_activation.loom | 1470 +++++++++++++++++ .../motifs/matmul/vector_decode_pair.loom | 188 +++ .../motifs/matmul/vector_postops.loom | 46 + .../motifs/matmul/vector_publish.loom | 21 + .../motifs/matmul/vector_publish_f32.loom | 22 + .../motifs/matmul/vector_q6_f32_body.loom | 748 +++++++++ .../loom-libs/ops/matmul/vector_f32_f32.loom | 126 ++ .../ops/matmul/vector_postops_f32_f32.loom | 177 ++ .../ops/matmul/vector_q6_f32_f32_kernel.loom | 39 + ...or_q6_selected_outputs_f32_f32_kernel.loom | 30 + tests/test-hrx-ops.cpp | 213 ++- 13 files changed, 4022 insertions(+), 233 deletions(-) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom 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 0fd957b3360c..d7ffd80bd75f 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 @@ -27,6 +27,16 @@ static constexpr KernelCatalogRef kMulMatAddF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatBiasAddF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_bias_add_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorBiasF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_residual_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorBiasAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_residual_f32_f32"); +static constexpr KernelCatalogRef kQuantizeQ8_1X4F32Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K64PlaneKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i4_k64_plane"); static constexpr KernelCatalogRef kMulMatSymmetricI4WmmaKernel = @@ -52,11 +62,11 @@ static constexpr KernelCatalogRef kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256"); static constexpr KernelCatalogRef kMulMatQ4KF16WmmaPrefillWave32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); -static constexpr KernelCatalogRef kMulMatQ6KPackedToken1F16WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_packed_token1_f16_wmma"); -static constexpr int64_t kMulMatQ6KPackedMaxOutputSize = 262144; -static constexpr KernelCatalogRef kMulMatQ6KI8PrepackedF16WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma"); +static constexpr KernelCatalogRef kMulMatVectorQ6F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q6_f32_f32"); +static constexpr int64_t kMulMatVectorDenseMaxOutputSize = 524288; +static constexpr int64_t kMulMatVectorPackedMaxOutputSize = 262144; +static constexpr int64_t kMulMatVectorSelectedOutputsMaxOutputSize = 262144; static constexpr KernelCatalogRef kMulMatQ6KF32WmmaPrefillWave32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_f32_wmma_prefill_wave32"); static constexpr KernelCatalogRef kMulMatQ6KF16WmmaPrefillWave32Kernel = @@ -70,8 +80,9 @@ static constexpr KernelCatalogRef kTopK8F32PartitionsRegisterKernel = static constexpr KernelCatalogRef kTopK128F32ReduceGatherRegisterKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_top_k128_f32_reduce_gather_register"); static constexpr KernelCatalogRef kFillNegativeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_fill_negative_f32"); -static constexpr KernelCatalogRef kMulMatQ6KPackedSelectedRefineToken1Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_packed_selected_refine_token1"); +static constexpr KernelCatalogRef kMulMatVectorQ6SelectedOutputsF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q6_selected_outputs_f32_f32"); +static constexpr int64_t kQwenVectorVocabularyCount = 151936; struct MulMatPostOpsMatch { const Value * input = nullptr; @@ -110,6 +121,20 @@ struct DecodeMulMatAddMatch { } }; +enum class VectorPublishFormat { + None, + F16, +}; + +struct VectorPublishPlan { + VectorPublishFormat fused_format = VectorPublishFormat::None; + ValueId fused_output = {}; + size_t fused_byte_count = 0; + bool publish_q8 = false; + ValueId q8_output = {}; + size_t q8_byte_count = 0; +}; + static bool is_bias_shape(const Value & value, int64_t output_size) { if (value.kind != ValueKind::External || value.type != GGML_TYPE_F32 || !value.contiguous || value.ne[0] != output_size) { @@ -260,6 +285,25 @@ static bool has_embedded_external_concat_projection_ancestor(const Graph & graph return false; } +static bool is_qwen_endpoint_rmsnorm_projection(const Graph & graph, const CommonMulMatMatch & match) { + if (!graph.has_index() || match.weight == nullptr || match.input == nullptr || match.weight->type != GGML_TYPE_Q6_K || + match.token_count != 1 || match.output_size != kQwenVectorVocabularyCount) { + return false; + } + + const GraphNode * mul = graph.index().producer(match.input->id); + if (mul == nullptr || mul->op != GGML_OP_MUL || mul->inputs.size() != 2) { + return false; + } + for (ValueId input : mul->inputs) { + const GraphNode * producer = graph.index().producer(input); + if (producer != nullptr && producer->op == GGML_OP_RMS_NORM) { + return true; + } + } + return false; +} + static bool match_symmetric_i4_low_row_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { // Disabled: this symmetric matmul route causes a substantial numeric performance regression. (void) context; @@ -369,17 +413,19 @@ static bool match_symmetric_i4_low_row_dispatch(const DispatchMatchContext & con return dispatch_match.status.success(); } -static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - constexpr int64_t selected_group_count = 96; - constexpr int64_t candidate_count = 128; - constexpr int64_t refine_tile_size = 64; - constexpr size_t partition_entry_count = 1024; +static bool match_q6_vector_selected_outputs_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + constexpr int64_t selected_group_count = 96; + constexpr int64_t selected_output_count = 128; + constexpr int64_t selected_output_chunk_size = 64; + constexpr size_t partition_entry_count = 1024; const CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KSymmetricI2ScanToken1Kernel, true); if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || match.input_size > 8192 || match.output_size < 65536 || match.output_size % 64 != 0 || + match.output_size > kMulMatVectorSelectedOutputsMaxOutputSize || match.output_size < 16 * match.input_size || !context.graph.index().consumers(match.output->id).empty() || !has_embedded_external_concat_projection_ancestor(context.graph, match.input->id, match.input_size)) { return false; @@ -398,8 +444,8 @@ static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & co const ValueId selected_groups(transient_base); const ValueId partial_values(transient_base + 1); const ValueId partial_ids(transient_base + 2); - const ValueId candidates(transient_base + 3); - const ValueId candidate_values(transient_base + 4); + const ValueId selected_outputs(transient_base + 3); + const ValueId selected_values(transient_base + 4); if (alternate == nullptr) { dispatch_match.transients.push_back( @@ -427,16 +473,20 @@ static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & co } } - dispatch_match.transients.push_back({ selected_groups, "common.mul_mat.q6_k_shortlist.selected_groups", + dispatch_match.transients.push_back({ selected_groups, "common.mul_mat.vector_q6_selected_outputs.groups", static_cast(selected_group_count) * sizeof(int32_t), 256 }); dispatch_match.transients.push_back( - { partial_values, "common.mul_mat.q6_k_shortlist.partial_values", partition_entry_count * sizeof(float), 256 }); + { partial_values, "common.mul_mat.vector_q6_selected_outputs.partial_values", + partition_entry_count * sizeof(float), 256 }); dispatch_match.transients.push_back( - { partial_ids, "common.mul_mat.q6_k_shortlist.partial_ids", partition_entry_count * sizeof(int32_t), 256 }); + { partial_ids, "common.mul_mat.vector_q6_selected_outputs.partial_ids", + partition_entry_count * sizeof(int32_t), 256 }); dispatch_match.transients.push_back( - { candidates, "common.mul_mat.q6_k_shortlist.candidates", candidate_count * sizeof(int32_t), 256 }); + { selected_outputs, "common.mul_mat.vector_q6_selected_outputs.indices", + selected_output_count * sizeof(int32_t), 256 }); dispatch_match.transients.push_back( - { candidate_values, "common.mul_mat.q6_k_shortlist.candidate_values", candidate_count * sizeof(float), 256 }); + { selected_values, "common.mul_mat.vector_q6_selected_outputs.values", + selected_output_count * sizeof(float), 256 }); Dispatch select_groups; select_groups.kernel = make_kernel_specialization(kSelectSymmetricI4K32GroupsKernel); @@ -485,8 +535,8 @@ static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & co reduce_top_k.kernel.integer_parameters.emplace("element_count", match.output_size); reduce_top_k.bindings.push_back({ partial_values, 0, partition_entry_count * sizeof(float) }); reduce_top_k.bindings.push_back({ partial_ids, 0, partition_entry_count * sizeof(int32_t) }); - reduce_top_k.bindings.push_back({ candidates, 0, candidate_count * sizeof(int32_t) }); - reduce_top_k.bindings.push_back({ candidate_values, 0, candidate_count * sizeof(float) }); + reduce_top_k.bindings.push_back({ selected_outputs, 0, selected_output_count * sizeof(int32_t) }); + reduce_top_k.bindings.push_back({ selected_values, 0, selected_output_count * sizeof(float) }); dispatch_match.dispatches.push_back(std::move(reduce_top_k)); Dispatch fill; @@ -495,49 +545,68 @@ static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & co fill.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch_match.dispatches.push_back(std::move(fill)); - for (int64_t candidate_offset = 0; candidate_offset < candidate_count; candidate_offset += refine_tile_size) { - Dispatch refine; - refine.kernel = make_kernel_specialization(kMulMatQ6KPackedSelectedRefineToken1Kernel); - refine.kernel.integer_parameters.emplace("token_count", match.token_count); - refine.kernel.integer_parameters.emplace("candidate_count", refine_tile_size); - refine.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size", - common_to_config_value(match.input_size)); - refine.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_size", - common_to_config_value(match.output_size)); - refine.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_accumulation", "0"); - refine.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.weight_offset", - common_to_config_value(static_cast(symmetric_i2_bytes))); - refine.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - refine.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, - kQ6KSymmetricI2PackedK256Row64ScaleRowLayout, match.weight->type, match.input_size, - match.output_size, match.weight->byte_count }); - refine.bindings.push_back({ candidates, static_cast(candidate_offset) * sizeof(int32_t), - static_cast(refine_tile_size) * sizeof(int32_t) }); - refine.bindings.push_back({ candidate_values, static_cast(candidate_offset) * sizeof(float), - static_cast(refine_tile_size) * sizeof(float) }); - refine.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - dispatch_match.dispatches.push_back(std::move(refine)); + for (int64_t selected_output_offset = 0; selected_output_offset < selected_output_count; + selected_output_offset += selected_output_chunk_size) { + Dispatch selected_output_pass; + selected_output_pass.kernel = make_kernel_specialization(kMulMatVectorQ6SelectedOutputsF32Kernel); + selected_output_pass.kernel.integer_parameters.emplace("token_count", match.token_count); + selected_output_pass.kernel.integer_parameters.emplace("selected_count", selected_output_chunk_size); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", + common_to_config_value(match.input_size)); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", + common_to_config_value(match.output_size)); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.weight_offset", + common_to_config_value( + static_cast(symmetric_i2_bytes))); + selected_output_pass.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + selected_output_pass.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, + kQ6KSymmetricI2PackedK256Row64ScaleRowLayout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + selected_output_pass.bindings.push_back( + { selected_outputs, static_cast(selected_output_offset) * sizeof(int32_t), + static_cast(selected_output_chunk_size) * sizeof(int32_t) }); + selected_output_pass.bindings.push_back( + { selected_values, static_cast(selected_output_offset) * sizeof(float), + static_cast(selected_output_chunk_size) * sizeof(float) }); + selected_output_pass.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch_match.dispatches.push_back(std::move(selected_output_pass)); } dispatch_match.covered_nodes.push_back(context.root_index); return dispatch_match.status.success(); } -static bool build_q6_k_aligned_skinny_dispatch(const CommonMulMatMatch & match, - size_t root_index, - KernelCatalogRef kernel, - DispatchMatch & dispatch_match) { +enum class Q6VectorStorage : int64_t { + Packed = 0, + I8 = 1, +}; + +static bool build_q6_vector_dispatch(const CommonMulMatMatch & match, + size_t root_index, + Q6VectorStorage storage, + size_t weight_bytes, + DispatchMatch & dispatch_match) { Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kernel); + dispatch.kernel = make_kernel_specialization(kMulMatVectorQ6F32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size", + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", common_to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_size", + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", common_to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.q6_storage", + common_to_config_value(static_cast(storage))); dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, kQ6KPackedK256Row64ScaleRowLayout, - match.weight->type, match.input_size, match.output_size, match.weight->byte_count }); + if (storage == Q6VectorStorage::I8) { + dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KI8K32Row64Layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KPackedK256Row64ScaleRowLayout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch_match.covered_nodes.push_back(root_index); @@ -545,57 +614,32 @@ static bool build_q6_k_aligned_skinny_dispatch(const CommonMulMatMatch & match, return true; } -static bool match_q6_k_aligned_skinny_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KPackedToken1F16WmmaKernel, false); - if (!match.matched()) { - match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KPackedToken1F16WmmaKernel, - true); - } +static bool build_q6_vector_dispatch(const CommonMulMatMatch & match, + size_t root_index, + DispatchMatch & dispatch_match) { + return build_q6_vector_dispatch(match, root_index, Q6VectorStorage::Packed, match.weight->byte_count, + dispatch_match); +} + +static bool match_q6_vector_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorQ6F32Kernel, true); if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || - match.weight->alias_source.value >= 0 || match.token_count > 16 || match.input_size % 256 != 0 || - match.output_size % 64 != 0 || match.output_size > kMulMatQ6KPackedMaxOutputSize || - (match.token_count > 1 && match.output_size < 8 * match.input_size) || + match.weight->alias_source.value >= 0 || match.token_count != 1 || match.input_size % 256 != 0 || + match.output_size % 64 != 0 || match.output_size > kMulMatVectorPackedMaxOutputSize || !context.graph.index().consumers(match.output->id).empty()) { return false; } - - return build_q6_k_aligned_skinny_dispatch(match, context.root_index, kMulMatQ6KPackedToken1F16WmmaKernel, - dispatch_match); -} - -static bool match_q6_k_i8_prepacked_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KI8PrepackedF16WmmaKernel, false); - if (!match.matched() || match.weight->type != GGML_TYPE_Q6_K || match.weight->alias_source.value >= 0 || - match.token_count <= 5 || match.token_count > 16 || match.input_size % 256 != 0 || match.output_size % 64 != 0 || - (match.output_size > match.input_size / 2 && match.token_count > 5)) { + if (is_qwen_endpoint_rmsnorm_projection(context.graph, match)) { return false; } - - const size_t block_count = static_cast(match.input_size / ggml_blck_size(GGML_TYPE_Q6_K)); - size_t materialized_weight_bytes = static_cast(match.output_size) * block_count; - if (materialized_weight_bytes > std::numeric_limits::max() / size_t{ 274 }) { + if (find_alternate_value(context.graph, context.plan, match.input->id, GGML_TYPE_Q8_1, + static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, + match.input_size)) != nullptr) { return false; } - materialized_weight_bytes *= size_t{ 274 }; - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatQ6KI8PrepackedF16WmmaKernel); - dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size", - common_to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_size", - common_to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_accumulation", "0"); - dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - dispatch.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, kQ6KI8K32Row64Layout, - match.weight->type, match.input_size, match.output_size, match.weight->byte_count }); - dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - dispatch_match.covered_nodes.push_back(context.root_index); - dispatch_match.dispatches.push_back(std::move(dispatch)); - return true; + return build_q6_vector_dispatch(match, context.root_index, dispatch_match); } static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, @@ -867,7 +911,8 @@ static bool match_q4_k_q8_1_x4_prefill_dispatch(const DispatchMatchContext & con match_packed_q8_1_x4_prefill_dispatch(context, dispatch_match); } -static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context) { +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, + bool allow_batch1_bias_residual) { MulMatPostOpsMatch match; CommonMulMatMatch root = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); @@ -878,15 +923,15 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con return match; } - const bool lowtoken_contraction = root.token_count <= 5 && root.weight->alias_source.value < 0 && - (root.weight_format == CommonMulMatWeightFormat::Q4K || - root.weight_format == CommonMulMatWeightFormat::Q6K) && - root.input_size >= 4096 && root.output_size >= 4096 && - root.output_size <= root.input_size && root.output_size % 64 == 0; - const bool token1_requires_q8_activation = root.token_count == 1 && !lowtoken_contraction; + const bool packed_row_contraction = root.token_count <= 5 && root.weight->alias_source.value < 0 && + (root.weight_format == CommonMulMatWeightFormat::Q4K || + root.weight_format == CommonMulMatWeightFormat::Q6K) && + root.input_size >= 4096 && root.output_size >= 4096 && + root.output_size <= root.input_size && root.output_size % 64 == 0; + const bool batch1_requires_q8_activation = root.token_count == 1 && !packed_row_contraction; const Value * current = root.output; - if (lowtoken_contraction) { + if (packed_row_contraction) { const GraphNode * reshape = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_RESHAPE); const Value * reshaped = reshape != nullptr ? common_graph_value(context.graph, reshape->output) : nullptr; if (reshaped != nullptr && is_layout_alias_node(context.graph, *reshape) && reshaped->contiguous && @@ -895,7 +940,7 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con current = reshaped; } } - const int add_limit = root.token_count == 1 || match.layout_node != nullptr ? 1 : 2; + const int add_limit = match.layout_node != nullptr || (root.token_count == 1 && !allow_batch1_bias_residual) ? 1 : 2; for (int add_index = 0; add_index < add_limit; ++add_index) { const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { @@ -962,10 +1007,14 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con match.input_size = root.input_size; match.output_size = root.output_size; match.token_count = root.token_count; - match.requires_q8_activation = token1_requires_q8_activation; + match.requires_q8_activation = batch1_requires_q8_activation; return match; } +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context) { + return match_mul_mat_postops(context, false); +} + static bool try_match_fused_unary(const DispatchMatchContext & context, CommonMulMatMatch & match) { if (!match.matched()) { return false; @@ -1080,28 +1129,462 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, return true; } -static bool match_q6_k_token1_final_projection_q8_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { +static bool prepare_vector_activation_binding(const DispatchMatchContext & context, + const Value & input, + const Value & weight, + int64_t input_size, + int64_t output_size, + int64_t token_count, + CommonMulMatWeightFormat weight_format, + DispatchMatch & dispatch_match, + DispatchBinding & activation, + CommonMulMatWeightFormat & kernel_weight_format, + const char *& weight_layout, + bool & uses_q8_activation) { + activation = { input.id, 0, input.byte_count }; + kernel_weight_format = weight_format; + weight_layout = nullptr; + uses_q8_activation = false; + + if (token_count != 1 || input_size % 256 != 0 || weight.alias_source.value >= 0 || output_size % 64 != 0) { + return true; + } + + const bool can_pack_q4 = weight_format == CommonMulMatWeightFormat::Q4K; + const bool can_pack_q6 = weight_format == CommonMulMatWeightFormat::Q6K; + if (!can_pack_q4 && !can_pack_q6) { + return true; + } + + DispatchBinding q8_activation = activation; + if (!common_prepare_q8_1_x4_input(context, input, input_size, token_count, dispatch_match, q8_activation, + CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + return true; + } + + activation = q8_activation; + kernel_weight_format = can_pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; + weight_layout = can_pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + uses_q8_activation = true; + return true; +} + +static bool is_k256_packed_row_layout(CommonMulMatWeightFormat format) { + return format == CommonMulMatWeightFormat::Q4KRow64 || format == CommonMulMatWeightFormat::Q6KRow64; +} + +static int64_t vector_max_output_size(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::Q4K: + case CommonMulMatWeightFormat::Q4KRow64: + case CommonMulMatWeightFormat::Q6K: + case CommonMulMatWeightFormat::Q6KRow64: + return kMulMatVectorPackedMaxOutputSize; + default: + return kMulMatVectorDenseMaxOutputSize; + } +} + +static bool vector_mul_mat_route_supports(const CommonMulMatMatch & match) { + return match.matched() && match.token_count == 1 && match.input_size >= 256 && + match.input_size <= 32768 && match.input_size % 32 == 0 && + match.output_size <= vector_max_output_size(match.weight_format); +} + +static bool vector_postops_route_supports(const MulMatPostOpsMatch & match) { + return match.matched() && match.token_count == 1 && match.input_size >= 256 && + match.input_size <= 32768 && match.input_size % 32 == 0 && + match.output_size <= vector_max_output_size(match.weight_format); +} + +static size_t vector_publish_byte_count(ggml_type type, int64_t token_count, int64_t output_size) { + if (token_count <= 0 || output_size <= 0) { + return 0; + } + return static_cast(token_count) * ggml_row_size(type, output_size); +} + +static bool is_vector_packed_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { + if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || + consumer->inputs[1] != input.id) { + return false; + } + + const Value * weight = common_graph_value(graph, consumer->inputs[0]); + const Value * output = common_graph_value(graph, consumer->output); + if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input.contiguous || !weight->contiguous || !output->contiguous || weight->alias_source.value >= 0) { + return false; + } + + const bool supported_weight = weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_Q6_K; + return supported_weight && input.ne[1] >= 1 && input.ne[1] <= 5 && input.ne[0] >= 256 && + input.ne[0] <= 32768 && input.ne[0] % 256 == 0 && input.ne[2] == 1 && input.ne[3] == 1 && + weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && weight->ne[1] <= 262144 && + weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && weight->ne[3] == 1 && + output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && + output->ne[3] == 1; +} + +static bool has_vector_packed_q8_consumer(const Graph & graph, const Value & value) { + if (!graph.has_index()) { + return false; + } + for (const GraphNode * consumer : graph.index().consumers(value.id)) { + if (is_vector_packed_q8_consumer(graph, consumer, value)) { + return true; + } + } + return false; +} + +static bool is_vector_f16_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { + if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || + consumer->inputs[1] != input.id) { + return false; + } + + const Value * weight = common_graph_value(graph, consumer->inputs[0]); + const Value * output = common_graph_value(graph, consumer->output); + if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input.contiguous || !weight->contiguous || !output->contiguous || weight->alias_source.value >= 0) { + return false; + } + + CommonMulMatWeightFormat format; + if (!common_mul_mat_format_for_type(weight->type, format)) { + return false; + } + return input.ne[0] == weight->ne[0] && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && + input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1 && + output->ne[2] == 1 && output->ne[3] == 1 && + common_mul_mat_uses_k16_major_f16(format, input.ne[0], output->ne[0], input.ne[1]); +} + +static bool has_vector_f16_consumer(const Graph & graph, const Value & value) { + if (!graph.has_index()) { + return false; + } + for (const GraphNode * consumer : graph.index().consumers(value.id)) { + if (is_vector_f16_consumer(graph, consumer, value)) { + return true; + } + } + return false; +} + +static VectorPublishPlan make_vector_publish_plan(const DispatchMatchContext & context, + const Value & output, + int64_t token_count, + int64_t output_size, + const DispatchMatch & dispatch_match) { + VectorPublishPlan plan; + if (has_vector_packed_q8_consumer(context.graph, output)) { + plan.publish_q8 = true; + plan.q8_output = ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + plan.q8_byte_count = vector_publish_byte_count(GGML_TYPE_Q8_1, token_count, output_size); + } else if (has_vector_f16_consumer(context.graph, output)) { + plan.fused_format = VectorPublishFormat::F16; + plan.fused_output = ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + plan.fused_byte_count = vector_publish_byte_count(GGML_TYPE_F16, token_count, output_size); + } + return plan; +} + +static int64_t vector_publish_format_config_value(VectorPublishFormat format) { + switch (format) { + case VectorPublishFormat::F16: + return GGML_TYPE_F16; + case VectorPublishFormat::None: + return 0; + } + return 0; +} + +static bool append_vector_q8_publish(const Value & output, + int64_t token_count, + int64_t output_size, + VectorPublishPlan & publish, + DispatchMatch & dispatch_match) { + if (!publish.publish_q8) { + return true; + } + if (publish.q8_byte_count == 0) { + return false; + } + + constexpr const char * name = "common.mul_mat.vector.q8_1_x4"; + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeQ8_1X4F32Kernel); + quantize.kernel.integer_parameters.emplace("token_count", token_count); + quantize.kernel.integer_parameters.emplace("input_size", output_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + common_to_config_value(token_count * output_size / 128)); + quantize.bindings.push_back({ output.id, 0, output.byte_count }); + quantize.bindings.push_back({ publish.q8_output, 0, publish.q8_byte_count }); + dispatch_match.dispatches.push_back(std::move(quantize)); + dispatch_match.transients.push_back({ publish.q8_output, name, publish.q8_byte_count, 256 }); + + Status status; + if (!dispatch_match.metadata.append_alternate_value( + { output.id, publish.q8_output, GGML_TYPE_Q8_1, publish.q8_byte_count, name }, status)) { + dispatch_match.status.append(status); + return false; + } + return true; +} + +static bool append_vector_f16_publish_metadata(const Value & output, + VectorPublishPlan & publish, + DispatchMatch & dispatch_match) { + if (publish.fused_format != VectorPublishFormat::F16) { + return true; + } + if (publish.fused_byte_count == 0) { + return false; + } + + constexpr const char * name = "common.mul_mat.vector.f16"; + dispatch_match.transients.push_back({ publish.fused_output, name, publish.fused_byte_count, 256 }); + Status status; + if (!dispatch_match.metadata.append_alternate_value( + { output.id, publish.fused_output, GGML_TYPE_F16, publish.fused_byte_count, name }, status)) { + dispatch_match.status.append(status); + return false; + } + return true; +} + +static DispatchBinding vector_next_output_binding(const Value & output, const VectorPublishPlan & publish) { + if (publish.fused_format == VectorPublishFormat::F16) { + return { publish.fused_output, 0, publish.fused_byte_count }; + } + return { output.id, 0, output.byte_count }; +} + +static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, + const CommonMulMatMatch & match, + DispatchMatch & dispatch_match) { + if (match.token_count != 1 || match.input_size % 32 != 0 || + match.output_size > vector_max_output_size(match.weight_format)) { + return false; + } + + if (match.input_size % 256 != 0 && is_k256_packed_row_layout(match.weight_format)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatVectorF32F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_unary_op", + std::to_string(unary_kind_config_value(match.output_unary_op))); + VectorPublishPlan publish = + make_vector_publish_plan(context, *match.output, match.token_count, match.output_size, dispatch_match); + if (!append_vector_f16_publish_metadata(*match.output, publish, dispatch_match)) { + return false; + } + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector.publish_format", + common_to_config_value(vector_publish_format_config_value(publish.fused_format))); + + DispatchBinding activation; + CommonMulMatWeightFormat weight_format = match.weight_format; + const char * weight_layout = nullptr; + bool uses_q8_activation = false; + if (!prepare_vector_activation_binding(context, *match.input, *match.weight, match.input_size, match.output_size, + match.token_count, match.weight_format, dispatch_match, activation, + weight_format, weight_layout, uses_q8_activation)) { + return false; + } + + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector.weight_format", + common_to_config_value(common_mul_mat_format_config_value(weight_format))); + if (uses_q8_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } + + dispatch.bindings.push_back(activation); + if (weight_layout != nullptr) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch.bindings.push_back(vector_next_output_binding(*match.output, publish)); + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.has_fused_unary) { + dispatch_match.covered_nodes.push_back(match.unary_node_index); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return append_vector_q8_publish(*match.output, match.token_count, match.output_size, publish, dispatch_match); +} + +static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & context, + const MulMatPostOpsMatch & match, + DispatchMatch & dispatch_match) { + if (match.token_count != 1 || match.input_size % 32 != 0 || + match.output_size > vector_max_output_size(match.weight_format)) { + return false; + } + + const bool has_tail = match.input_size % 256 != 0; + if (has_tail && is_k256_packed_row_layout(match.weight_format)) { + return false; + } + + KernelCatalogRef kernel = {}; + if (match.has_bias && match.has_residual) { + kernel = kMulMatVectorBiasAddF32F32Kernel; + } else if (match.has_bias) { + kernel = kMulMatVectorBiasF32F32Kernel; + } else if (match.has_residual) { + kernel = kMulMatVectorAddF32F32Kernel; + } else { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.apply_bias", + common_to_config_value(match.has_bias ? int64_t{ 1 } : int64_t{ 0 })); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.apply_residual", + common_to_config_value(match.has_residual ? int64_t{ 1 } : int64_t{ 0 })); + VectorPublishPlan publish = + make_vector_publish_plan(context, *match.residual_output, match.token_count, match.output_size, dispatch_match); + if (!append_vector_f16_publish_metadata(*match.residual_output, publish, dispatch_match)) { + return false; + } + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector.publish_format", + common_to_config_value(vector_publish_format_config_value(publish.fused_format))); + + DispatchBinding activation; + CommonMulMatWeightFormat weight_format = match.weight_format; + const char * weight_layout = nullptr; + bool uses_q8_activation = false; + if (!prepare_vector_activation_binding(context, *match.input, *match.weight, match.input_size, match.output_size, + match.token_count, match.weight_format, dispatch_match, activation, + weight_format, weight_layout, uses_q8_activation)) { + return false; + } + + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector_postops.weight_format", + common_to_config_value(common_mul_mat_format_config_value(weight_format))); + if (uses_q8_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } + + dispatch.bindings.push_back(activation); + if (weight_layout != nullptr) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + if (match.has_bias) { + dispatch.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + } + if (match.has_residual) { + dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); + } + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + dispatch.bindings.push_back(vector_next_output_binding(*match.residual_output, publish)); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes)) { + return false; + } + if (match.layout_node != nullptr && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.layout_node, + dispatch_match.covered_nodes)) { + return false; + } + for (const GraphNode * add_node : match.add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return append_vector_q8_publish(*match.residual_output, match.token_count, match.output_size, publish, + dispatch_match); +} + +static bool match_vector_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorF32F32Kernel, true); + if (!match.matched() || match.token_count != 1) { + return false; + } + return build_vector_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_vector_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorF32F32Kernel, true); + if (!match.matched() || match.token_count != 1 || !try_match_fused_unary(context, match)) { + return false; + } + return build_vector_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_vector_mul_mat_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = match_mul_mat_postops(context, true); + if (!match.matched() || match.token_count != 1) { + return false; + } + return build_vector_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static bool match_q6_vector_final_projection_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, true); if (!match.matched()) { match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); } if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || - match.token_count < 1 || match.token_count > 5 || match.weight->alias_source.value >= 0 || + match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || match.output_size % 64 != 0 || - match.output_size > kMulMatQ6KPackedMaxOutputSize || + match.output_size > kMulMatVectorPackedMaxOutputSize || !context.graph.index().consumers(match.output->id).empty()) { return false; } - if (match.token_count == 1 && match.output_size >= 16 * match.input_size) { - return build_q6_k_aligned_skinny_dispatch(match, context.root_index, kMulMatQ6KPackedToken1F16WmmaKernel, - dispatch_match); + if (is_qwen_endpoint_rmsnorm_projection(context.graph, match)) { + return false; } - return build_mul_mat_dispatch(context, match, dispatch_match, - CommonQ8ActivationPolicy::AllowStandaloneQuantize); + if (match.output_size < 16 * match.input_size) { + return false; + } + + return build_q6_vector_dispatch(match, context.root_index, dispatch_match); } static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { @@ -1109,6 +1592,9 @@ static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, if (!match.matched()) { return false; } + if (vector_postops_route_supports(match)) { + return false; + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(match.kernel); @@ -1304,6 +1790,9 @@ static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, if (!match.matched()) { return false; } + if (vector_mul_mat_route_supports(match)) { + return false; + } if ((match.weight_format == CommonMulMatWeightFormat::Q4K || match.weight_format == CommonMulMatWeightFormat::Q6K) && match.output_size % 64 == 0 && common_is_supported_dense_output_size(match.output_size) && match.weight->alias_source.value < 0) { @@ -1320,17 +1809,32 @@ static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & conte if (!match.matched()) { return false; } + MulMatPostOpsMatch postops; + postops.input = match.root.input; + postops.weight = match.root.weight; + postops.projection_output = match.root.output; + postops.residual_input = match.residual_input; + postops.residual_output = match.residual_output; + postops.kernel = kMulMatVectorAddF32F32Kernel; + postops.weight_format = match.root.weight_format; + postops.input_size = match.root.input_size; + postops.output_size = match.root.output_size; + postops.token_count = match.root.token_count; + postops.has_residual = true; + if (vector_postops_route_supports(postops)) { + return false; + } return build_decode_mul_mat_add_dispatch(match, context, dispatch_match); } void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { registry.add({ - "common.mul_mat.q6_k_token1_shortlist", + "common.mul_mat.vector_q6_selected_outputs", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 310, DispatchSource::Common, - match_q6_k_token1_shortlist_dispatch, + match_q6_vector_selected_outputs_dispatch, }); registry.add({ "common.mul_mat.q6_k_prefill_wave32", @@ -1341,20 +1845,20 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_q6_k_prefill_wave32_dispatch, }); registry.add({ - "common.mul_mat.q6_k_token1_final_projection_q8", + "common.mul_mat.vector_q6_final_projection_q8", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 305, DispatchSource::Common, - match_q6_k_token1_final_projection_q8_dispatch, + match_q6_vector_final_projection_q8_dispatch, }); registry.add({ - "common.mul_mat.q6_k_aligned_skinny", + "common.mul_mat.vector_q6", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 300, DispatchSource::Common, - match_q6_k_aligned_skinny_dispatch, + match_q6_vector_dispatch, }); registry.add({ "common.mul_mat.q5_k_iq4_xs_symmetric_i4_lowrow", @@ -1364,14 +1868,6 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_symmetric_i4_low_row_dispatch, }); - registry.add({ - "common.mul_mat.q6_k_i8_prepacked", - GGML_OP_MUL_MAT, - DispatchMatchKind::Fused, - 285, - DispatchSource::Common, - match_q6_k_i8_prepacked_dispatch, - }); registry.add({ "common.mul_mat.q5_k_iq4_xs_q8_1_x4_prefill", GGML_OP_MUL_MAT, @@ -1404,6 +1900,22 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_q5_k_symmetric_i8_prefill_dispatch, }); + registry.add({ + "common.mul_mat_postops.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_vector_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat_unary.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_vector_mul_mat_unary_dispatch, + }); registry.add({ "common.mul_mat_unary.f32_f32_wmma", GGML_OP_MUL_MAT, @@ -1412,6 +1924,14 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_mul_mat_unary_dispatch, }); + registry.add({ + "common.mul_mat.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 190, + DispatchSource::Common, + match_vector_mul_mat_dispatch, + }); registry.add({ "common.mul_mat_postops.f32_f32_wmma", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 2f14648c4913..60e5d4557cef 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -32,6 +32,18 @@ { "path": "ops/mul_mat_f32_f32_wmma.loom" }, + { + "path": "ops/matmul/vector_f32_f32.loom" + }, + { + "path": "ops/matmul/vector_postops_f32_f32.loom" + }, + { + "path": "motifs/matmul/vector_publish.loom" + }, + { + "path": "motifs/matmul/vector_decode_pair.loom" + }, { "path": "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom" }, @@ -116,6 +128,12 @@ { "path": "ops/mul_mat_q6_k_packed_f16_wmma.loom" }, + { + "path": "ops/matmul/vector_q6_f32_f32_kernel.loom" + }, + { + "path": "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom" + }, { "path": "ops/mul_mat_dual_q4_f32_decode.loom" }, @@ -173,6 +191,18 @@ { "path": "motifs/mul_mat_f32_f32_wmma_core.loom" }, + { + "path": "motifs/matmul/vector_activation.loom" + }, + { + "path": "motifs/matmul/vector_postops.loom" + }, + { + "path": "motifs/matmul/vector_q6_f32_body.loom" + }, + { + "path": "motifs/matmul/vector_publish_f32.loom" + }, { "path": "motifs/mul_mat_quantized_f16_prefill.loom" }, @@ -1488,6 +1518,274 @@ "motifs/binary_f32_apply.loom" ] }, + { + "name": "ggml_mul_mat_vector_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_f32_f32", + "source": "ops/matmul/vector_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_bias_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_bias_f32_f32", + "source": "ops/matmul/vector_postops_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_postops_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_residual_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_residual_f32_f32", + "source": "ops/matmul/vector_postops_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual_input", + "residual_output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_postops_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_bias_residual_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_bias_residual_f32_f32", + "source": "ops/matmul/vector_postops_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_postops_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, { "name": "ggml_mul_mat_f32_f32_narrow_split_k4", "family": "loom_libs", @@ -4234,57 +4532,10 @@ ] }, { - "name": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", - "family": "loom_libs", - "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", - "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_q6_k_packed_f16_wmma.loom" - ], - "library_sources": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" - ] - }, - "compile_dependencies": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" - ] - }, - { - "name": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "name": "ggml_mul_mat_vector_q6_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", - "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "symbol": "ggml_mul_mat_vector_q6_f32_f32", + "source": "ops/matmul/vector_q6_f32_f32_kernel.loom", "workload_parameters": [ { "name": "token_count", @@ -4311,20 +4562,14 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_q6_k_packed_f16_wmma.loom" + "ops/matmul/vector_q6_f32_f32_kernel.loom" ], "library_sources": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" + "motifs/matmul/vector_q6_f32_body.loom" ] }, "compile_dependencies": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" + "motifs/matmul/vector_q6_f32_body.loom" ] }, { @@ -4470,8 +4715,8 @@ "bindings": [ "partial_values", "partial_ids", - "candidate_output", - "value_output" + "selected_outputs", + "selected_values" ], "binding_access": [ "read", @@ -4543,17 +4788,17 @@ ] }, { - "name": "ggml_mul_mat_q6_k_packed_selected_refine_token1", + "name": "ggml_mul_mat_vector_q6_selected_outputs_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_q6_k_packed_selected_refine_token1", - "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "symbol": "ggml_mul_mat_vector_q6_selected_outputs_f32_f32", + "source": "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom", "workload_parameters": [ { "name": "token_count", "type": "index" }, { - "name": "candidate_count", + "name": "selected_count", "type": "index" } ], @@ -4563,15 +4808,15 @@ "type": "index" }, { - "name": "candidate_count", + "name": "selected_count", "type": "index" } ], "bindings": [ "input", "weight", - "candidates", - "exact_output", + "selected_outputs", + "selected_values", "output" ], "binding_access": [ @@ -4585,20 +4830,16 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_q6_k_packed_f16_wmma.loom" + "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom" ], "library_sources": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" + "motifs/matmul/vector_q6_f32_body.loom", + "motifs/matmul/vector_publish_f32.loom" ] }, "compile_dependencies": [ - "motifs/q6_k_f16.loom", - "motifs/dequant.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" + "motifs/matmul/vector_q6_f32_body.loom", + "motifs/matmul/vector_publish_f32.loom" ] }, { @@ -5721,5 +5962,7 @@ } ], "link_modules": [], - "plan_cases": [] + "plan_cases": [], + "sources": [], + "kernels": [] } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom new file mode 100644 index 000000000000..60e0a5b34e53 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom @@ -0,0 +1,1470 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.matmul.vector.activation_format = 0 : index +config.def @ggml.matmul.vector.publish_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +// Shared vector-shaped matmul reduction pipeline. +// +// Op wrappers provide publication and output-group hooks by defining: +// ggml.matmul.vector_activation.publish_vector4 +// ggml.matmul.vector_activation.finish_output_group +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: i1, %arg6: f32, %arg7: f32, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +func.decl @ggml_matmul_vector_f32_decode_pair(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + + +func.def pure inline @ggml_matmul_vector_q8_geometry(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %q4_format = index.constant 44 : index + %q6_format = index.constant 46 : index + %q8_format = index.constant 9 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %native_q4 = index.cmp eq, %weight_format, %q4_format : index + %native_q6 = index.cmp eq, %weight_format, %q6_format : index + %q8_input = index.cmp eq, %activation_format, %q8_format : index + %q4_min_blocks = index.constant 327680 : index + %blocks = index.div %input_size, %c256 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large_payload = index.cmp uge, %weight_blocks, %q4_min_blocks : index + %q4_can_stage = scalar.andi %native_q4, %large_payload : i1 + %native_staged = scalar.ori %q4_can_stage, %native_q6 : i1 + %native_q8 = scalar.andi %native_staged, %q8_input : i1 + %q4_min_k = index.constant 16384 : index + %q6_min_k = index.constant 8192 : index + %min_k = scf.select %native_q4, %q4_min_k, %q6_min_k : index + %min_n = index.constant 4096 : index + %five_tokens = index.cmp eq, %token_capacity, %c5 : index + %long_k = index.cmp uge, %input_size, %min_k : index + %wide_n = index.cmp uge, %output_size, %min_n : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large = scalar.andi %long_k, %wide_n : i1 + %long_reduction = scalar.andi %large, %contracts : i1 + %staged_shape = scalar.andi %five_tokens, %long_reduction : i1 + %staged = scalar.andi %native_q8, %staged_shape : i1 + %rows = scf.select %native_q4, %c64, %c32 : index + %result = scf.select %staged, %rows, %c0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %block_bytes = index.constant 144 : index + %half_cache_bytes = index.constant 16777216 : index + %cache_bytes = index.constant 33554432 : index + %payload_bytes = index.mul %weight_blocks, %block_bytes : index + %large_half = index.cmp uge, %payload_bytes, %half_cache_bytes : index + %fits_cache = index.cmp ule, %payload_bytes, %cache_bytes : index + %cache_range = scalar.andi %large_half, %fits_cache : i1 + %remainder = index.rem %blocks, %c2 : index + %even_blocks = index.cmp eq, %remainder, %c0 : index + %reduces_rows = index.cmp ult, %output_size, %input_size : index + %native_q4_q8 = scalar.andi %native_q4, %q8_input : i1 + %plain = index.cmp eq, %result, %c0 : index + %q4_plain = scalar.andi %native_q4_q8, %plain : i1 + %five_wide = scalar.andi %five_tokens, %wide_n : i1 + %reduction = scalar.andi %reduces_rows, %even_blocks : i1 + %partition_shape = scalar.andi %five_wide, %reduction : i1 + %partition_payload = scalar.andi %partition_shape, %cache_range : i1 + %partitioned = scalar.andi %q4_plain, %partition_payload : i1 + %k_partitions = scf.select %partitioned, %c2, %c1 : index + func.return %result, %k_partitions : index, index +} + +func.def pure inline @ggml_matmul_vector_uses_dot(%token_capacity: index, %input_size: index) -> (i1) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %few_tokens = index.cmp ule, %token_capacity, %c5 : index + %tail = index.rem %input_size, %c32 : index + %aligned_reduction = index.cmp eq, %tail, %c0 : index + %use_dot = scalar.andi %few_tokens, %aligned_reduction : i1 + func.return %use_dot : i1 +} + +func.def pure inline @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity: index, %input_size: index, %publish_format: index) -> (i1) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %one_token = index.cmp eq, %token_capacity, %c1 : index + %tail = index.rem %input_size, %c32 : index + %aligned_reduction = index.cmp eq, %tail, %c0 : index + %plain_publish = index.cmp eq, %publish_format, %c0 : index + %shape = scalar.andi %one_token, %aligned_reduction : i1 + %f32_shape = scalar.andi %shape, %f32_input : i1 + %use_pair = scalar.andi %f32_shape, %plain_publish : i1 + func.return %use_pair : i1 +} + +template.def<@ggml.matmul.vector_activation.launch> pure @ggml_matmul_vector_activation_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_groups = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_groups = index.div %padded_token_count, %c32 : index + %c5 = index.constant 5 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %vector_reduction = func.call pure @ggml_matmul_vector_uses_dot(%token_capacity, %input_size) : (index, index) -> (i1) + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %small_output = index.cmp ult, %output_size, %c64 : index + %narrow = scalar.andi %vector_reduction, %small_output : i1 + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%token_capacity, %input_size, %output_size, %weight_format) : (index, index, index, index) -> (index, index) + %staged = index.cmp ugt, %staged_rows, %c0 : index + %staged_rounding = index.sub %staged_rows, %c1 : index + %staged_threads = index.mul %staged_rows, %c16 : index + %cohort_rows = index.div %c8, %k_partitions : index + %cohort_rounding = index.sub %cohort_rows, %c1 : index + %base_rows_per_group = scf.select %narrow, %c4, %cohort_rows : index + %c1024 = index.constant 1024 : index + %c4096 = index.constant 4096 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %batched_tokens = index.cmp ugt, %token_capacity, %c1 : index + %long_reduction = index.cmp uge, %input_size, %c4096 : index + %partition_tail = index.rem %input_size, %c1024 : index + %full_partitions = index.cmp eq, %partition_tail, %c0 : index + %batch_f32 = scalar.andi %batched_tokens, %f32_input : i1 + %partitionable = scalar.andi %long_reduction, %full_partitions : i1 + %partition_rows = scalar.andi %batch_f32, %partitionable : i1 + %narrow_threads = scf.select %partition_rows, %c1024, %c256 : index + %base_workgroup_size = scf.select %narrow, %narrow_threads, %c128 : index + %base_rounding = scf.select %narrow, %c3, %cohort_rounding : index + %row_rounding = scf.select %staged, %staged_rounding, %base_rounding : index + %rows_per_group = scf.select %staged, %staged_rows, %base_rows_per_group : index + %workgroup_size = scf.select %staged, %staged_threads, %base_workgroup_size : index + %rounded_n = index.add %output_size, %row_rounding : index + %direct_row_groups = index.div %rounded_n, %rows_per_group : index + %publish_format = config.get @ggml.matmul.vector.publish_format : index + %pair_decode = func.call pure @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity, %input_size, %publish_format) : (index, index, index) -> (i1) + %padded_pair_output = index.add %output_size, %c1 : index + %pair_groups = index.div %padded_pair_output, %c2 : index + %pair_workgroup_size = index.constant 64 : index + %reduction_launch_x = scf.select %pair_decode, %pair_groups, %direct_row_groups : index + %reduction_workgroup_size = scf.select %pair_decode, %pair_workgroup_size, %workgroup_size : index + %launch_x = scf.select %vector_reduction, %reduction_launch_x, %output_groups : index + %launch_y = scf.select %vector_reduction, %c1, %token_groups : index + template.return %launch_x, %launch_y, %c1, %reduction_workgroup_size : index, index, index, index +} + +func.def inline @ggml_matmul_vector_combine_k_partitions(%value: f32, %partitions: index) -> (f32) { + %c1 = index.constant 1 : index + %partitioned = index.cmp ugt, %partitions, %c1 : index + %result = scf.if %partitioned -> (f32) { + %offset = scalar.constant 16 : i32 + %width = scalar.constant 32 : i32 + %peer, %ok = kernel.subgroup.shuffle %value, %offset, %width : f32, i32, i32 + %sum = scalar.addf %value, %peer : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_matmul_vector_reduce_cohort_f32(%value: f32) -> (f32) { + %width = scalar.constant 32 : i32 + %off1 = scalar.constant 1 : i32 + %peer1, %ok1 = kernel.subgroup.shuffle %value, %off1, %width : f32, i32, i32 + %sum1 = scalar.addf %value, %peer1 : f32 + %off2 = scalar.constant 2 : i32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %off2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %off4 = scalar.constant 4 : i32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %off4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %off8 = scalar.constant 8 : i32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %off8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + func.return %sum8 : f32 +} + +// Narrow vector projections assign a full wave to each row. +// Both cohort sizes use the existing wrapper publication hooks. +func.def inline @ggml_matmul_vector_reduce_f32(%values: vector<4xf32>, %full_wave: i1) -> (f32) { + %zero = scalar.constant 0.0 : f32 + %width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %sum = vector.reduce %values, %zero : vector<4xf32>, f32 + %peer1, %ok1 = kernel.subgroup.shuffle %sum, %xor1, %width : f32, i32, i32 + %sum1 = scalar.addf %sum, %peer1 : f32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %xor2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %xor4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %xor8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + %result = scf.if %full_wave -> (f32) { + %xor16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %xor16, %width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer16 : f32 + %lane32 = scalar.constant 32 : i32 + %peer32 = kernel.subgroup.broadcast %sum16 from %lane32 : f32, i32 + %sum64 = scalar.addf %sum16, %peer32 : f32 + scf.yield %sum64 : f32 + } else { + scf.yield %sum8 : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_matmul_vector_f32_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %tokens = index.assume %token_count [range(%token_count, 1, 1)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %workgroup_size = kernel.workgroup.size : index + %one_wave_per_row = index.cmp eq, %workgroup_size, %c256 : index + %c1024 = index.constant 1024 : index + %partition_group = index.cmp eq, %workgroup_size, %c1024 : index + %narrow_output = index.cmp ult, %n, %c64 : index + %partitioned = scalar.andi %partition_group, %narrow_output : i1 + %full_wave = scalar.ori %one_wave_per_row, %partitioned : i1 + %k_partitions = scf.select %partitioned, %c4, %c1 : index + %cohort_width = scf.select %full_wave, %c64, %c16 : index + %rows_per_group = scf.select %full_wave, %c4, %c8 : index + %wave_rows = scf.select %full_wave, %c1, %c4 : index + %step_width = scf.select %full_wave, %c256, %c64 : index + %steps_per_block = scf.select %full_wave, %c1, %c4 : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %lane16 = index.rem %lane, %cohort_width : index + %cohort = index.div %lane, %cohort_width : index + %row_group_base = index.mul %channel_group, %rows_per_group : index + %row_subgroup = index.div %subgroup, %k_partitions : index + %k_partition = index.rem %subgroup, %k_partitions : index + %wave_add = index.mul %row_subgroup, %wave_rows : index + %wave_base = index.add %row_group_base, %wave_add : index + %row = index.add %wave_base, %cohort : index + %valid_row = index.cmp ult, %row, %n : index + %packet = index.rem %lane16, %c8 : index + %lane_group = index.div %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %c255 = index.constant 255 : index + %padded_k = index.add %k, %c255 : index + %blocks = index.div %padded_k, %c256 : index + %step_rounding = index.sub %step_width, %c1 : index + %padded_steps = index.add %k, %step_rounding : index + %steps = index.div %padded_steps, %step_width : index + %partition_steps = index.div %steps, %k_partitions : index + %step_begin = index.mul %k_partition, %partition_steps : index + %step_end = index.add %step_begin, %partition_steps : index + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %k) : (index, index) -> (offset) + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %a_na, %w_na = buffer.assume.noalias %input, %weight : buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%blocks]x256xf32> + %total0, %total1, %total2, %total3, %total4 = scf.for %step = [%step_begin to %step_end step %c1](%acc0 = %zero : vector<4xf32>, %acc1 = %zero : vector<4xf32>, %acc2 = %zero : vector<4xf32>, %acc3 = %zero : vector<4xf32>, %acc4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block = index.div %step, %steps_per_block : index + %quarter = index.rem %step, %steps_per_block : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %step_k = index.mul %step, %step_width : index + %kk = index.add %step_k, %lane_k : index + %valid_k = index.cmp ult, %kk, %k : index + %valid_row_k = scalar.andi %valid_row, %valid_k : i1 + %local_step_k = index.mul %quarter, %step_width : index + %local_k = index.add %local_step_k, %lane_k : index + %input_mask = vector.mask.range [%kk to %k step %c1] : index -> vector<4xi1> + %ww = scf.if %valid_row_k -> (vector<4xf32>) { + %format44 = index.constant 44 : index + %packed = index.cmp eq, %weight_format, %format44 : index + %wh = scf.if %packed -> (vector<4xf32>) { + %v = func.call @ggml_q4k_native_row64_f32_vector4(%w_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %v = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %w_na, %row_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } + scf.yield %wh : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next0 = scf.if %valid0 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t0, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc0 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc0 : vector<4xf32> + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next1 = scf.if %valid1 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t1, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc1 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc1 : vector<4xf32> + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next2 = scf.if %valid2 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t2, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc2 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc2 : vector<4xf32> + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next3 = scf.if %valid3 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t3, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc3 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc3 : vector<4xf32> + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next4 = scf.if %valid4 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t4, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc4 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc4 : vector<4xf32> + } + scf.yield %next0, %next1, %next2, %next3, %next4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %c5 = index.constant 5 : index + %stage_width = scf.select %partitioned, %c16, %c8 : index + %stage_elements = index.mul %stage_width, %c5 : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%stage_width]xf32> + %local_row0 = index.add %wave_add, %cohort : index + %partial_offset = index.mul %k_partition, %c4 : index + %local_row = index.add %local_row0, %partial_offset : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total0, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total1, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total2, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total3, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total4, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %rows_per_group, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_group_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, 4), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %stage_width : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_origin = scf.if %partitioned -> (index) { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 64), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } else { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 36), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values0 = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %values = scf.if %partitioned -> (vector<4xf32>) { + %stage_origin1 = index.add %stage_origin, %c4 : index + %stage_origin2 = index.add %stage_origin1, %c4 : index + %stage_origin3 = index.add %stage_origin2, %c4 : index + %values1 = vector.load %flat_stage[%stage_origin1] : view<[%stage_elements]xf32> -> vector<4xf32> + %values2 = vector.load %flat_stage[%stage_origin2] : view<[%stage_elements]xf32> -> vector<4xf32> + %values3 = vector.load %flat_stage[%stage_origin3] : view<[%stage_elements]xf32> -> vector<4xf32> + %sum01 = vector.addf %values0, %values1 : vector<4xf32> + %sum012 = vector.addf %sum01, %values2 : vector<4xf32> + %sum = vector.addf %sum012, %values3 : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values0 : vector<4xf32> + } + %publish_format = config.get @ggml.matmul.vector.publish_format : index + template.apply<@ggml.matmul.vector_activation.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %publish_format, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %row_rounding = index.sub %rows_per_group, %c1 : index + %rounded_n = index.add %n, %row_rounding : index + %output_groups = index.div %rounded_n, %rows_per_group : index + template.apply<@ggml.matmul.vector_activation.finish_output_group>(%tokens, %n, %output_groups, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Copy the original Q8_1_x4 records for up to 2048 reduction elements. +// Call uniformly at a slab boundary; only the final slab may be partial. +func.def inline @ggml_matmul_vector_q8_1_x4_stage_slab(%input: buffer, %stage: buffer, %capacity: index, %slab_begin: index, %blocks: index, %global_row_bytes: offset) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c144 = index.constant 144 : index + %c576 = index.constant 576 : index + %c18 = index.constant 18 : index + %base = index.constant 0 : offset + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %workgroup_size = kernel.workgroup.size : index + %thread = kernel.workitem.id : index + %stage_packets = index.mul %capacity, %c144 : index + %stage_elements = index.mul %capacity, %c576 : index + %stage_view = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xi32> + %slab_limit = index.add %slab_begin, %c8 : index + %full_slab = index.cmp ule, %slab_limit, %blocks : index + %slab_end = scf.select %full_slab, %slab_limit, %blocks : index + %active_blocks0 = index.sub %slab_end, %slab_begin : index + %active_blocks = index.assume %active_blocks0 [range(%active_blocks0, 1, 8)] : index + %active_packets = index.mul %active_blocks, %c18 : index + %slab_byte_add = index.scale %slab_begin, %block_bytes : index, offset -> offset + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %copy_packet = [%thread to %stage_packets step %workgroup_size] { + %copy_token = index.div %copy_packet, %c144 : index + %copy_inner = index.rem %copy_packet, %c144 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + scf.if %copy_valid { + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %copy_bytes = index.scale %copy_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slab_byte_add : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + %copy_origin = index.mul %copy_packet, %c4 : index + vector.store %values, %stage_view[%copy_origin] : vector<4xi32>, view<[%stage_elements]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_matmul_vector_q8_publish(%tokens: index, %n: index, %rows_per_group: index, %row_base: index, %local_row: index, %block_lane: index, %channel_group: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %publish_format: index, %total_g0: f32, %total_g1: f32, %total_g2: f32, %total_g3: f32, %total_g4: f32, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %leader = index.cmp eq, %block_lane, %c0 : index + %width = scalar.constant 32 : i32 + %stage_rows = index.constant 5 : index + %stage_elements = index.mul %stage_rows, %rows_per_group : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%rows_per_group]xf32> + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum = scalar.addf %total_g0, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum = scalar.addf %total_g1, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum = scalar.addf %total_g2, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum = scalar.addf %total_g3, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum = scalar.addf %total_g4, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %rows_per_group, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %row_limit = index.sub %rows_per_group, %c4 : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, %row_limit), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %rows_per_group : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_limit = index.sub %stage_elements, %c4 : index + %stage_origin = index.assume %stage_origin0 [range(%stage_origin0, 0, %stage_limit), mul(%stage_origin0, 4)] : index + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %publish_format = config.get @ggml.matmul.vector.publish_format : index + template.apply<@ggml.matmul.vector_activation.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %publish_format, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %rounding = index.sub %rows_per_group, %c1 : index + %rounded_n = index.add %n, %rounding : index + %output_groups = index.div %rounded_n, %rows_per_group : index + template.apply<@ggml.matmul.vector_activation.finish_output_group>(%tokens, %n, %output_groups, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_q4_prefetch_native_words(%weight: buffer, %row_group_block: index, %block: index, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<2xi32>) { + %zero = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %group_block = index.add %row_group_block, %block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %header_view = buffer.view %weight[%wb] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%wc] : buffer -> view<8x64x4xi32> + %headers = vector.load %header_view[%zero] : view<4xi32> -> vector<4xi32> + %raw = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %headers, %raw : vector<4xi32>, vector<2xi32> +} + +// Native Q4_K bytes and Q8_1_x4 activations, with one 16-lane cohort per row. +// Retain the wrapper publication hooks so bias, residual and RMSNorm fusions +// share this contraction without a separate dispatch or intermediate output. +func.def pure inline @ggml_matmul_vector_q4_q8_token_pairs(%capacity: index, %tokens: index, %k: index, %n: index) -> (i1) { + %c4 = index.constant 4 : index + %c4096 = index.constant 4096 : index + %capacity4 = index.cmp eq, %capacity, %c4 : index + %tokens4 = index.cmp eq, %tokens, %c4 : index + %four_tokens = scalar.andi %capacity4, %tokens4 : i1 + %long_input = index.cmp uge, %k, %c4096 : index + %wide_output = index.cmp uge, %n, %c4096 : index + %contracts = index.cmp ule, %n, %k : index + %wide = scalar.andi %long_input, %wide_output : i1 + %contracting = scalar.andi %wide, %contracts : i1 + %eligible = scalar.andi %four_tokens, %contracting : i1 + func.return %eligible : i1 +} + +func.def inline @ggml_matmul_vector_q4_q8_adjacent_scale_min(%scale0: i32, %scale1: i32, %scale2: i32, %pair0: index) -> (i32, i32, i32, i32) { + %pair = index.assume %pair0 [range(%pair0, 0, 3)] : index + %two = index.constant 2 : index + %sixteen = index.constant 16 : index + %low = index.cmp ult, %pair, %two : index + %pair_lane = index.rem %pair, %two : index + %shift_index = index.mul %pair_lane, %sixteen : index + %shift = index.cast %shift_index : index to i32 + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %eight_i32 = scalar.constant 8 : i32 + %low_mask = scalar.constant 3855 : i32 + %high_mask = scalar.constant 12336 : i32 + %byte_mask = scalar.constant 255 : i32 + %high_shift = scalar.addi %shift, %two_i32 : i32 + %min_shift = scalar.addi %shift, %four_i32 : i32 + %s_source = scf.select %low, %scale0, %scale2 : i32 + %m_source = scf.select %low, %scale1, %scale2 : i32 + %s_high_shift = scf.select %low, %shift, %high_shift : i32 + %m_low_shift = scf.select %low, %shift, %min_shift : i32 + %s_lo_raw = scalar.shrui %s_source, %shift : i32 + %s_lo = scalar.andi %s_lo_raw, %low_mask : i32 + %s_hi_raw = scalar.shrui %scale0, %s_high_shift : i32 + %s_hi = scalar.andi %s_hi_raw, %high_mask : i32 + %s = scalar.ori %s_lo, %s_hi : i32 + %m_lo_raw = scalar.shrui %m_source, %m_low_shift : i32 + %m_lo = scalar.andi %m_lo_raw, %low_mask : i32 + %m_hi_raw = scalar.shrui %scale1, %s_high_shift : i32 + %m_hi = scalar.andi %m_hi_raw, %high_mask : i32 + %m = scalar.ori %m_lo, %m_hi : i32 + %s0 = scalar.andi %s, %byte_mask : i32 + %m0 = scalar.andi %m, %byte_mask : i32 + %s1 = scalar.shrui %s, %eight_i32 : i32 + %m1 = scalar.shrui %m, %eight_i32 : i32 + func.return %s0, %m0, %s1, %m1 : i32, i32, i32, i32 +} + +func.def inline @ggml_matmul_vector_q4_q8_dot8(%paired_tokens: i1, %weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %value = scf.if %paired_tokens -> (f32) { + %zero = vector.constant 0 : vector<1xi32> + %wi0 = vector.extract %weight[0] : vector<2xi32> -> i32 + %ai0 = vector.extract %activation[0] : vector<2xi32> -> i32 + %ww0 = vector.splat %wi0 : vector<1xi32> + %aw0 = vector.splat %ai0 : vector<1xi32> + %w0 = vector.bitcast %ww0 : vector<1xi32> to vector<4xi8> + %a0 = vector.bitcast %aw0 : vector<1xi32> to vector<4xi8> + %d0 = vector.dot4i %w0, %a0, %zero : vector<4xi8>, vector<4xi8>, vector<1xi32> + %wi1 = vector.extract %weight[1] : vector<2xi32> -> i32 + %ai1 = vector.extract %activation[1] : vector<2xi32> -> i32 + %ww1 = vector.splat %wi1 : vector<1xi32> + %aw1 = vector.splat %ai1 : vector<1xi32> + %w1 = vector.bitcast %ww1 : vector<1xi32> to vector<4xi8> + %a1 = vector.bitcast %aw1 : vector<1xi32> to vector<4xi8> + %d1 = vector.dot4i %w1, %a1, %d0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %sum = vector.extract %d1[0] : vector<1xi32> -> i32 + %result = scalar.sitofp %sum : i32 to f32 + scf.yield %result : f32 + } else { + %original = func.call @ggml_dot_u8_s8_vector8_f32(%weight, %activation) : (vector<2xi32>, vector<2xi32>) -> (f32) + scf.yield %original : f32 + } + func.return %value : f32 +} + +func.def inline @ggml_matmul_vector_q4_q8_1_x4_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %paired_tokens = func.call pure @ggml_matmul_vector_q4_q8_token_pairs(%capacity, %tokens, %k, %n) : (index, index, index, index) -> (i1) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = index.add %channel_group, %c0 : index + %c64_rows = index.constant 64 : index + %direct_rows = index.div %c8, %k_partitions : index + %rows_per_group = scf.select %staged, %c64_rows, %direct_rows : index + %row_base = index.mul %wg, %rows_per_group : index + %wave_rows = index.div %c4, %k_partitions : index + %wave_row = index.mul %wave, %wave_rows : index + %cohort_width = index.mul %c16, %k_partitions : index + %row_half = index.div %lane, %cohort_width : index + %cohort_lane = index.rem %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %k_partitions : index + %k_team = index.div %cohort_lane, %c16 : index + %k_begin = index.mul %k_team, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %global_q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %slab_row_bytes = index.constant 2304 : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %channel_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = index.constant 8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_header = vector.constant 0 : vector<4xi32> + %zero_raw = vector.constant 0 : vector<2xi32> + %prime_header, %prime_raw = scf.if %staged -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %c0, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %zero_header, %zero_raw : vector<4xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_header, %last_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_header = %prime_header : vector<4xi32>, %carry_raw = %prime_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32>) { + %block_in_slab = index.rem %block_base, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %new_slab = scalar.andi %staged, %slab_boundary : i1 + scf.if %new_slab { + func.call @ggml_matmul_vector_q8_1_x4_stage_slab(%input_na, %q8_stage, %capacity, %block_base, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %block = index.add %block_base, %k_begin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_block, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %staged -> (vector<4xi32>) { + scf.yield %carry_header : vector<4xi32> + } else { + %header = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %header : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %pair_scale0, %pair_min0, %pair_scale1, %pair_min1 = scf.if %paired_tokens -> (i32, i32, i32, i32) { + %s0, %m0, %s1, %m1 = func.call @ggml_matmul_vector_q4_q8_adjacent_scale_min(%g_s0, %g_s1, %g_s2, %pair) : (i32, i32, i32, index) -> (i32, i32, i32, i32) + scf.yield %s0, %m0, %s1, %m1 : i32, i32, i32, i32 + } else { + %zero_i32 = scalar.constant 0 : i32 + scf.yield %zero_i32, %zero_i32, %zero_i32, %zero_i32 : i32, i32, i32, i32 + } + %g_scale0i, %g_min0i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale0, %pair_min0 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale1, %pair_min1 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %staged -> (vector<2xi32>) { + scf.yield %carry_raw : vector<2xi32> + } else { + %raw = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %raw : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_header, %next_raw = scf.if %has_next -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %next_block, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %carry_header, %carry_raw : vector<4xi32>, vector<2xi32> + } + %input_block = scf.select %staged, %block_in_slab, %safe_block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0 = scf.if %valid0 -> (f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ad0h = view.load %ds_view[%q8_meta0] : view<8xf16> -> f16 + %ad1h = view.load %ds_view[%q8_meta1] : view<8xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = view.load %ds_view[%q8_sum0] : view<8xf16> -> f16 + %as1h = view.load %ds_view[%q8_sum1] : view<8xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1 = scf.if %valid1 -> (f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + scf.if %paired_tokens { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2 = scf.if %valid2 -> (f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3 = scf.if %valid3 -> (f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4 = scf.if %valid4 -> (f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_header, %next_raw : f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32> + } + %combined0 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g0, %k_partitions) : (f32, index) -> (f32) + %combined1 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g1, %k_partitions) : (f32, index) -> (f32) + %combined2 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g2, %k_partitions) : (f32, index) -> (f32) + %combined3 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g3, %k_partitions) : (f32, index) -> (f32) + %combined4 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g4, %k_partitions) : (f32, index) -> (f32) + %publish_format = config.get @ggml.matmul.vector.publish_format : index + func.call @ggml_matmul_vector_q8_publish(%tokens, %n, %rows_per_group, %row_base, %local_row, %cohort_lane, %channel_group, %output_accumulation, %output_unary_op, %epsilon, %publish_format, %combined0, %combined1, %combined2, %combined3, %combined4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, index, index, f32, index, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Native Q6_K carrier with signed I8 dot products and F32 scaling. +func.def inline @ggml_matmul_vector_q6_prefetch_native_words(%weight_na: buffer, %row_group_block: index, %block: index, %row_lane: index, %half16: index, %group: index, %half128: index, %ql_side: index, %packet: index) -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %c2 = index.constant 2 : index + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %group_block = index.add %row_group_block, %block : index + %block_base = index.scale %group_block, %record_bytes : index, offset -> offset + %scale_base = index.add %block_base, %scale_add : offset + %raw_base = index.add %block_base, %raw_add : offset + %d_view = buffer.view %weight_na[%block_base] : buffer -> view<64xf16> + %s_view = buffer.view %weight_na[%scale_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_na[%raw_base] : buffer -> view<2x3x64x8xi32> + %dh = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %si = view.load %s_view[%row_lane, %half16, %group] : view<64x2x8xi8> -> i8 + %ql = vector.load %raw_view[%half128, %ql_side, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + %qh = vector.load %raw_view[%half128, %c2, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + func.return %dh, %si, %ql, %qh : f16, i8, vector<4xi32>, vector<4xi32> +} + +func.def inline @ggml_matmul_vector_q6_q8_activation_values(%pipelined: i1, %token: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index) -> (vector<16xi8>, f32) { + %code_offset = index.constant 16 : offset + %a, %ad = scf.if %pipelined -> (vector<16xi8>, f32) { + %base = index.constant 0 : offset + %c4 = index.constant 4 : index + %q8_block = index.assume %q8_block0 [range(%q8_block0, 0, 7)] : index + %word = index.add %q8_word0, %c4 : index + %q_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x36xi32> + %ds_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x72xf16> + %words = vector.load %q_view[%current_slot, %token, %q8_block, %word] : view<2x[%capacity]x8x36xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%current_slot, %token, %q8_block, %q8_meta] : view<2x[%capacity]x8x72xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } else { + %tb = index.scale %token, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %words = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%q8_meta] : view<8xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } + func.return %a, %ad : vector<16xi8>, f32 +} + +func.def inline @ggml_matmul_vector_q6_q8_dot16(%pipelined: i1, %w: vector<16xi8>, %a: vector<16xi8>) -> (i32) { + %sum = scf.if %pipelined -> (i32) { + %di0 = vector.constant 0 : vector<1xi32> + %dw0 = vector.slice %w[0] : vector<16xi8> -> vector<4xi8> + %da0 = vector.slice %a[0] : vector<16xi8> -> vector<4xi8> + %di1 = vector.dot4i %dw0, %da0, %di0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw1 = vector.slice %w[4] : vector<16xi8> -> vector<4xi8> + %da1 = vector.slice %a[4] : vector<16xi8> -> vector<4xi8> + %di2 = vector.dot4i %dw1, %da1, %di1 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw2 = vector.slice %w[8] : vector<16xi8> -> vector<4xi8> + %da2 = vector.slice %a[8] : vector<16xi8> -> vector<4xi8> + %di3 = vector.dot4i %dw2, %da2, %di2 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw3 = vector.slice %w[12] : vector<16xi8> -> vector<4xi8> + %da3 = vector.slice %a[12] : vector<16xi8> -> vector<4xi8> + %di4 = vector.dot4i %dw3, %da3, %di3 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %chain_sum = vector.extract %di4[0] : vector<1xi32> -> i32 + scf.yield %chain_sum : i32 + } else { + %init = vector.constant 0 : vector<4xi32> + %dot = vector.dot4i %w, %a, %init : vector<16xi8>, vector<16xi8>, vector<4xi32> + %seed = scalar.constant 0 : i32 + %direct_sum = vector.reduce %dot, %seed : vector<4xi32>, i32 + scf.yield %direct_sum : i32 + } + func.return %sum : i32 +} + +func.def inline @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input: buffer, %capacity0: index, %slot_begin: index, %blocks: index, %global_row_bytes: offset) -> (vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c18 = index.constant 18 : index + %c72 = index.constant 72 : index + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %copy_packet = index.rem %thread, %total_packets : index + %copy_token = index.div %copy_packet, %c72 : index + %copy_inner = index.rem %copy_packet, %c72 : index + %remaining = index.sub %blocks, %slot_begin : index + %full = index.cmp uge, %remaining, %c4 : index + %active0 = scf.select %full, %c4, %remaining : index + %active = index.assume %active0 [range(%active0, 1, 4)] : index + %active_packets = index.mul %active, %c18 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + %safe_inner = scf.select %copy_valid, %copy_inner, %c0 : index + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %slot_bytes = index.scale %slot_begin, %block_bytes : index, offset -> offset + %copy_bytes = index.scale %safe_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slot_bytes : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + func.return %values : vector<4xi32> +} +func.def inline @ggml_matmul_vector_q8_1_x4_publish_half_slab(%stage: buffer, %capacity0: index, %slot0: index, %values: vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %slot = index.assume %slot0 [range(%slot0, 0, 1)] : index + %c4 = index.constant 4 : index + %c72 = index.constant 72 : index + %base = index.constant 0 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %thread_valid = index.cmp ult, %thread, %total_packets : index + scf.if %thread_valid { + %copy_token0 = index.div %thread, %c72 : index + %copy_token = index.assume %copy_token0 [range(%copy_token0, 0, 4), lt(%copy_token0, %capacity)] : index + %copy_inner = index.rem %thread, %c72 : index + %origin = index.mul %copy_inner, %c4 : index + %destination = buffer.view %stage[%base] : buffer -> view<2x[%capacity]x288xi32> + vector.store %values, %destination[%slot, %copy_token, %origin] : vector<4xi32>, view<2x[%capacity]x288xi32> + } + func.return +} + + +func.def inline @ggml_matmul_vector_q6_q8_accumulate(%pipelined: i1, %tokens: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index, %w: vector<16xi8>, %ws: f32, %acc_g0: f32, %acc_g1: f32, %acc_g2: f32, %acc_g3: f32, %acc_g4: f32) -> (f32, f32, f32, f32, f32) { + %code_offset = index.constant 16 : offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %computed_g0 = scf.if %valid0 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t0, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g0 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %computed_g1 = scf.if %valid1 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t1, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g1 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %computed_g2 = scf.if %valid2 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t2, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g2 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %computed_g3 = scf.if %valid3 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t3, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g3 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %computed_g4 = scf.if %valid4 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t4, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g4 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + func.return %computed_g0, %computed_g1, %computed_g2, %computed_g3, %computed_g4 : f32, f32, f32, f32, f32 +} + +func.def inline @ggml_matmul_vector_q6_q8_1_x4_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %rows_per_group = scf.select %staged, %c32, %c8 : index + %row_base = index.mul %channel_group, %rows_per_group : index + %wave_row = index.mul %wave, %c4 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %block_lane = index.rem %lane, %c16 : index + %group = index.div %block_lane, %c2 : index + %half16 = index.rem %block_lane, %c2 : index + %half128 = index.div %group, %c4 : index + %group_in_half = index.rem %group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift0 = index.mul %nibble, %c4 : index + %nibble_shift1 = index.cast %nibble_shift0 : index to i32 + %nibble_shift = vector.splat %nibble_shift1 : vector<4xi32> + %high_shift0 = index.mul %group_in_half, %c2 : index + %high_shift1 = index.cast %high_shift0 : index to i32 + %high_shift = vector.splat %high_shift1 : vector<4xi32> + %packet = index.mul %half16, %c4 : index + %q8_group_words = index.mul %group_in_half, %c8 : index + %q8_word0 = index.add %q8_group_words, %packet : index + %q8_meta = index.mul %group_in_half, %c2 : index + %blocks = index.div %k, %c256 : index + %q8_groups = index.div %k, %c128 : index + %q8_block_bytes = index.constant 144 : offset + %global_q8_row_bytes = index.scale %q8_groups, %q8_block_bytes : index, offset -> offset + %native_block_bytes = index.constant 210 : index + %cache_bytes = index.constant 33554432 : index + %weight_blocks = index.mul %blocks, %n : index + %weight_bytes = index.mul %weight_blocks, %native_block_bytes : index + %exceeds_cache = index.cmp uge, %weight_bytes, %cache_bytes : index + %pipelined = scalar.andi %staged, %exceeds_cache : i1 + %full_slab_row_bytes = index.constant 2304 : offset + %half_slab_row_bytes = index.constant 1152 : offset + %slab_row_bytes = scf.select %pipelined, %half_slab_row_bytes, %full_slab_row_bytes : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %code_offset = index.constant 16 : offset + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %staged_row_group = index.div %channel_group, %c2 : index + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %staged_row_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %input_na, %weight_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = scf.select %pipelined, %c4, %c8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_half = scalar.constant 0.0 : f16 + %zero_code = scalar.constant 0 : i8 + %zero_words = vector.constant 0 : vector<4xi32> + %prime_d, %prime_s, %prime_ql, %prime_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %pd, %ps, %pql, %pqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %c0, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %pd, %ps, %pql, %pqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %zero_half, %zero_code, %zero_words, %zero_words : f16, i8, vector<4xi32>, vector<4xi32> + } + scf.if %pipelined { + %prime_q8 = func.call @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input_na, %capacity, %c0, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + func.call @ggml_matmul_vector_q8_1_x4_publish_half_slab(%q8_stage, %capacity, %c0, %prime_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_d, %last_s, %last_ql, %last_qh = scf.for %block = [%c0 to %blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_d = %prime_d : f16, %carry_s = %prime_s : i8, %carry_ql = %prime_ql : vector<4xi32>, %carry_qh = %prime_qh : vector<4xi32>) -> (f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32>) { + %block_in_slab = index.rem %block, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %direct_stage = scalar.andi %staged, %slab_boundary : i1 + %true = scalar.constant true : i1 + %not_pipelined = scalar.xori %pipelined, %true : i1 + %stage_current = scalar.andi %direct_stage, %not_pipelined : i1 + scf.if %stage_current { + func.call @ggml_matmul_vector_q8_1_x4_stage_slab(%input_na, %q8_stage, %capacity, %block, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %slab_number = index.div %block, %slab_blocks : index + %current_slot = index.rem %slab_number, %c2 : index + %following_slab = index.add %slab_number, %c1 : index + %following_begin = index.mul %following_slab, %slab_blocks : index + %following_exists = index.cmp ult, %following_begin, %blocks : index + %next_slot = index.rem %following_slab, %c2 : index + %staged_future = scalar.andi %pipelined, %following_exists : i1 + %current_d, %current_s, %current_ql, %current_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + %wd, %ws8, %wql, %wqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } + %d = scalar.extf %current_d : f16 to f32 + %sf = scalar.sitofp %current_s : i8 to f32 + %ws = scalar.mulf %d, %sf : f32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %four = vector.constant 4 : vector<4xi32> + %ql_shifted = vector.shrui %current_ql, %nibble_shift : vector<4xi32> + %low = vector.andi %ql_shifted, %nibble_mask : vector<4xi32> + %qh_shifted = vector.shrui %current_qh, %high_shift : vector<4xi32> + %high0 = vector.andi %qh_shifted, %high_mask : vector<4xi32> + %high = vector.shli %high0, %four : vector<4xi32> + %code = vector.ori %low, %high : vector<4xi32> + // Guard bits prevent subtraction from borrowing across packed bytes. + %guard_mask = vector.constant -2139062144 : vector<4xi32> + %center = vector.constant 538976288 : vector<4xi32> + %guarded = vector.ori %code, %guard_mask : vector<4xi32> + %centered = vector.subi %guarded, %center : vector<4xi32> + %signed_code = vector.xori %centered, %guard_mask : vector<4xi32> + %w = vector.bitcast %signed_code : vector<4xi32> to vector<16xi8> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_d, %next_s, %next_ql, %next_qh = scf.if %has_next -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %wd, %ws8, %wql, %wqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %next_block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } + %input_block = scf.select %staged, %block_in_slab, %block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %half128 : index + %q8_local_block_add = index.scale %q8_block, %q8_block_bytes : index, offset -> offset + %slot_row = index.mul %current_slot, %capacity : index + %slot_byte_add0 = index.scale %slot_row, %slab_row_bytes : index, offset -> offset + %slot_byte_add = scf.select %pipelined, %slot_byte_add0, %base : offset + %q8_block_add = index.add %slot_byte_add, %q8_local_block_add : offset + %slot_complete = index.cmp eq, %block_in_slab, %c3 : index + %publish_next = scalar.andi %staged_future, %slot_complete : i1 + %next_g0, %next_g1, %next_g2, %next_g3, %next_g4 = scf.if %publish_next -> (f32, f32, f32, f32, f32) { + %next_q8 = func.call @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input_na, %capacity, %following_begin, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_matmul_vector_q6_q8_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + func.call @ggml_matmul_vector_q8_1_x4_publish_half_slab(%q8_stage, %capacity, %next_slot, %next_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } else { + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_matmul_vector_q6_q8_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_d, %next_s, %next_ql, %next_qh : f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32> + } + %publish_format = config.get @ggml.matmul.vector.publish_format : index + func.call @ggml_matmul_vector_q8_publish(%tokens, %n, %rows_per_group, %row_base, %local_row, %block_lane, %channel_group, %output_accumulation, %output_unary_op, %epsilon, %publish_format, %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, index, index, f32, index, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + + +template.def<@ggml.matmul.vector_activation.core> device @ggml_matmul_vector_activation_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %q8_format = index.constant 9 : index + %uses_q8 = index.cmp eq, %activation_format, %q8_format : index + scf.if %uses_q8 { + %q6_format = index.constant 46 : index + %uses_q6 = index.cmp eq, %weight_format, %q6_format : index + scf.if %uses_q6 { + func.call @ggml_matmul_vector_q6_q8_1_x4_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_vector_q4_q8_1_x4_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + %token_capacity = config.get @ggml.workload.token_capacity : index + %publish_format = config.get @ggml.matmul.vector.publish_format : index + %pair_decode = func.call pure @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity, %input_size0, %publish_format) : (index, index, index) -> (i1) + scf.if %pair_decode { + func.call @ggml_matmul_vector_f32_decode_pair(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %publish_format, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_vector_f32_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom new file mode 100644 index 000000000000..1e0d4200a532 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom @@ -0,0 +1,188 @@ +template.decl @ggml.matmul.vector_activation.publish_pair2(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: i1, %arg6: f32, %arg7: f32, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +func.def inline @ggml_matmul_vector_decode_pair_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block0 = scf.select %valid_block, %block, %c0 : index + %safe_block, %launch_block_count = index.assume %safe_block0, %block_count [lt(%safe_block0, %block_count)] : index, index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32> + %vector_half_base = index.mul %vector_half, %c128 : index + %vector_offset = index.mul %vector_index, %c4 : index + %input_index0 = index.add %vector_half_base, %vector_offset : index + %input_index1 = index.add %input_index0, %c32 : index + %input_index2 = index.add %input_index0, %c64 : index + %input_index3 = index.add %input_index0, %c96 : index + %block_k_base = index.mul %safe_block, %c256 : index + %k0 = index.add %block_k_base, %input_index0 : index + %k1 = index.add %block_k_base, %input_index1 : index + %k2 = index.add %block_k_base, %input_index2 : index + %k3 = index.add %block_k_base, %input_index3 : index + %mask0 = vector.mask.range [%k0 to %input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %input_size step %c1] : index -> vector<4xi1> + %mask2 = vector.mask.range [%k2 to %input_size step %c1] : index -> vector<4xi1> + %mask3 = vector.mask.range [%k3 to %input_size step %c1] : index -> vector<4xi1> + %input0 = vector.load.mask %input_view[%token, %safe_block, %input_index0], %mask0, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input1 = vector.load.mask %input_view[%token, %safe_block, %input_index1], %mask1, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input2 = vector.load.mask %input_view[%token, %safe_block, %input_index2], %mask2, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input3 = vector.load.mask %input_view[%token, %safe_block, %input_index3], %mask3, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +func.def inline @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %input_size: index, %row: index, %block: index, %lane: index, %weight: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = vector.constant 0.0 : vector<4xf32> + %zero_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %packet = index.rem %itid, %c8 : index + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %group_base = index.mul %vector_half, %c4 : index + %group0 = index.add %group_base, %c0 : index + %group1 = index.add %group_base, %c1 : index + %group2 = index.add %group_base, %c2 : index + %group3 = index.add %group_base, %c3 : index + %block_k_base = index.mul %safe_block, %c256 : index + %half_base = index.mul %vector_half, %c128 : index + %packet_k = index.mul %packet, %c4 : index + %k_half = index.add %block_k_base, %half_base : index + %k0 = index.add %k_half, %packet_k : index + %k1 = index.add %k0, %c32 : index + %k2 = index.add %k0, %c64 : index + %k3 = index.add %k0, %c96 : index + %valid_k0 = index.cmp ult, %k0, %input_size : index + %valid_k1 = index.cmp ult, %k1, %input_size : index + %valid_k2 = index.cmp ult, %k2, %input_size : index + %valid_k3 = index.cmp ult, %k3, %input_size : index + %weight0 = scf.if %valid_k0 -> (vector<4xf32>) { + %decoded = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group0, %packet, %k0) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight1 = scf.if %valid_k1 -> (vector<4xf32>) { + %decoded = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group1, %packet, %k1) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight2 = scf.if %valid_k2 -> (vector<4xf32>) { + %decoded = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group2, %packet, %k2) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight3 = scf.if %valid_k3 -> (vector<4xf32>) { + %decoded = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group3, %packet, %k3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %mul0 = vector.mulf %input0, %weight0 : vector<4xf32> + %mul1 = vector.mulf %input1, %weight1 : vector<4xf32> + %mul2 = vector.mulf %input2, %weight2 : vector<4xf32> + %mul3 = vector.mulf %input3, %weight3 : vector<4xf32> + %sum01 = vector.addf %mul0, %mul1 : vector<4xf32> + %sum23 = vector.addf %mul2, %mul3 : vector<4xf32> + %sum = vector.addf %sum01, %sum23 : vector<4xf32> + %block_contribution = vector.reduce %sum, %zero_f32 : vector<4xf32>, f32 + %contribution = scf.select %valid_block, %block_contribution, %zero_f32 : f32 + func.return %contribution : f32 +} + +func.def inline @ggml_matmul_vector_f32_decode_pair(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %tokens = index.assume %token_count [range(%token_count, 1, 1)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %lane0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = scalar.constant 0.0 : f32 + %token_valid = index.cmp ult, %token_group, %tokens : index + %safe_token0 = scf.select %token_valid, %token_group, %c0 : index + %token, %launch_tokens = index.assume %safe_token0, %tokens [lt(%safe_token0, %tokens)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %channel_group, %c2 : index + %row_valid = index.cmp ult, %row00, %n : index + %safe_row0 = scf.select %row_valid, %row00, %c0 : index + %row0, %launch_n = index.assume %safe_row0, %n [lt(%safe_row0, %n)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_n : index + %publish_output = scalar.andi %token_valid, %row_valid : i1 + %padded_k = index.add %k, %c255 : index + %block_count = index.div %padded_k, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %zero : f32, %row_acc1 = %zero : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_matmul_vector_decode_pair_load_f32_block(%launch_tokens, %k, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %k, %row0, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + %row1_contribution = func.call @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %k, %row1, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %zero : f32 + } + %next0 = scalar.addf %row_acc0, %contribution0 : f32 + %next1 = scalar.addf %row_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %sum0 = kernel.workgroup.reduce %acc0 : f32 + %sum1 = kernel.workgroup.reduce %acc1 : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %publish_pair = scalar.andi %publish_output, %is_lane_zero : i1 + template.apply<@ggml.matmul.vector_activation.publish_pair2>(%publish_pair, %token, %row0, %launch_tokens, %launch_n, %row1_valid, %sum0, %sum1, %output_accumulation, %output_unary_op, %publish_format, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, i1, f32, f32, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + func.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom new file mode 100644 index 000000000000..5b782b08a182 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom @@ -0,0 +1,46 @@ +template.decl @ggml.matmul.vector_activation.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.vector_activation.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.vector_activation.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.vector_activation.add_bias_vector4> device @ggml_matmul_vector_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.vector_activation.add_residual_vector4> device @ggml_matmul_vector_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.vector_activation.store_vector4> device @ggml_matmul_vector_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom new file mode 100644 index 000000000000..6008d253a25d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom @@ -0,0 +1,21 @@ +// Shared publication helpers for vector-shaped matmul outputs. +template.decl @ggml.matmul.vector_publish.next_vector4(%publish_word: i1, %publish_format: index, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %next_output: buffer) + +template.def<@ggml.matmul.vector_publish.next_vector4> device @ggml_matmul_vector_publish_next_vector4(%publish_word: i1, %publish_format: index, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %f16_format = index.constant 1 : index + %base = index.constant 0 : offset + %publishes_f16 = index.cmp eq, %publish_format, %f16_format : index + %publishes_next = scalar.andi %publish_word, %publishes_f16 : i1 + scf.if %publishes_next { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %next_output_noalias = buffer.assume.noalias %next_output : buffer + %next_output_view = buffer.view %next_output_noalias[%base] : buffer -> view<[%token_count]x[%output_size]xf16> + %f16_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store.mask %f16_values, %next_output_view[%token, %channel], %mask : vector<4xf16>, view<[%token_count]x[%output_size]xf16>, vector<4xi1> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom new file mode 100644 index 000000000000..64738a8ff0a8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom @@ -0,0 +1,22 @@ +// F32 publication stages shared across vector implementations. +template.decl @ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) + +template.def<@ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32> device @ggml_matmul_vector_publish_f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) { + %base = index.constant 0 : offset + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %publishes = index.cmp ult, %tid, %selected_count : index + %selected_noalias, %values_noalias, %output_noalias = buffer.assume.noalias %selected_outputs, %values, %output : buffer, buffer, buffer + %selected_view = buffer.view %selected_noalias[%base] : buffer -> view<64xi32> + %values_view = buffer.view %values_noalias[%base] : buffer -> view<64xf32> + scf.if %publishes { + %output_i32 = view.load %selected_view[%tid] : view<64xi32> -> i32 + %output0 = index.cast %output_i32 : i32 to index + %output1 = index.assume %output0 [range(%output0, 0, 262143)] : index + %selected_output, %bounded_output_size = index.assume %output1, %output_size [lt(%output1, %output_size)] : index, index + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%bounded_output_size]xf32> + %value = view.load %values_view[%tid] : view<64xf32> -> f32 + view.store %value, %output_view[%selected_output] : f32, view<[%bounded_output_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom new file mode 100644 index 000000000000..93d341238477 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom @@ -0,0 +1,748 @@ +// Shared Q6-weight vector matmul reductions. +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector.weight_offset : %value: index where [range(%value, 0, 2147483647), mul(%value, 256)] + +template.decl @ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32(%result: vector<16xf16>, %token_offset: index, %token_start: index, %bounded_token_count: index, %bounded_output_size: index, %channel_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) + +template.decl @ggml.matmul.vector_q6_f32.store_dense_outputs_wave32(%result: vector<8xf32>, %token_offset: index, %token_start: index, %token_count: index, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) + +template.decl @ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32(%result: vector<8xf32>, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) + +template.def<@ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32> device @ggml_matmul_vector_q6_f32_store_fragment_outputs_wave32(%result: vector<16xf16>, %token_offset: index, %token_start: index, %bounded_token_count: index, %bounded_output_size: index, %channel_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) { + %lane = kernel.subgroup.lane.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %publish_token0 = index.div %lane, %c2 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c2 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index + %publish_channel_add = index.mul %publish_packet, %c8 : index + %token_base = index.add %token_start, %token_offset : index + %token = index.add %token_base, %publish_token : index + %subgroup_channel_base = index.add %channel_base, %subgroup_channel_add : index + %channel = index.add %subgroup_channel_base, %publish_channel_add : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %writes = scalar.andi %valid_token, %valid_channel : i1 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + vector.fragment.store %result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes { + %bounded_token, %published_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%published_token_count]x[%bounded_output_size]xf32> + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %wide = vector.extf %values : vector<8xf16> to vector<8xf32> + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> + %published = scf.if %accumulates_output -> (vector<8xf32>) { + %residual = vector.load.mask %bounded_output_view[%bounded_token, %channel], %mask, %c0_f32x8 : view<[%published_token_count]x[%bounded_output_size]xf32>, vector<8xi1>, vector<8xf32> + %sum = vector.addf %residual, %wide : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %wide : vector<8xf32> + } + vector.store.mask %published, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%published_token_count]x[%bounded_output_size]xf32>, vector<8xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32> device @ggml_matmul_vector_q6_f32_store_single_row_overwrite_wave32(%result: vector<8xf32>, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %channel = index.add %channel_base, %subgroup_channel_add : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_channel_base = index.add %channel, %lane_channel : index + %stores_token0 = index.cmp eq, %lane_token, %c0 : index + scf.if %stores_token0 { + %dst = buffer.view %output[%base] : buffer -> view<1x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + view.store %value, %dst[%c0, %bounded_channel] : f32, view<1x[%output_size]xf32> + } + } + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32> device @ggml_matmul_vector_q6_f32_store_dense_outputs_wave32(%result: vector<8xf32>, %token_offset: index, %token_start: index, %token_count: index, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) { + %base = index.constant 0 : offset + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %token = index.add %token_start, %token_offset : index + %channel = index.add %channel_base, %subgroup_channel_add : index + %token_end = index.add %token, %c16 : index + %channel_end = index.add %channel, %c16 : index + %valid_token = index.cmp ule, %token_end, %token_count : index + %valid_channel = index.cmp ule, %channel_end, %output_size : index + %valid = scalar.andi %valid_token, %valid_channel : i1 + scf.if %valid { + %layout = encoding.layout.strided [1, %output_size] : encoding + %dst = buffer.view %output[%base] : buffer -> view<[%output_size]x[%token_count]xf32, %layout> + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + %published = scf.if %accumulates -> (vector<8xf32>) { + %old = vector.fragment.load %dst[%channel, %token] shape [%c16, %c16] : view<[%output_size]x[%token_count]xf32, %layout> -> vector<8xf32> + %sum = vector.addf %old, %result : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %result : vector<8xf32> + } + vector.fragment.store %published, %dst[%channel, %token] shape [%c16, %c16] : vector<8xf32>, view<[%output_size]x[%token_count]xf32, %layout> + } else { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_token = index.add %token, %lane_token : index + %partial_channel_base = index.add %channel, %lane_channel : index + %token_valid = index.cmp ult, %partial_token, %token_count : index + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + scf.if %token_valid { + %bounded_token, %published_token_count = index.assume %partial_token, %token_count [lt(%partial_token, %token_count)] : index, index + %dst = buffer.view %output[%base] : buffer -> view<[%published_token_count]x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + %published = scf.if %accumulates -> (f32) { + %old = view.load %dst[%bounded_token, %bounded_channel] : view<[%published_token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %value, %old : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + view.store %published, %dst[%bounded_token, %bounded_channel] : f32, view<[%published_token_count]x[%output_size]xf32> + } + } + } + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.body_packed_storage> device @ggml_matmul_vector_q6_f32_body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_base = index.mul %output_block, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_base = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<8xf32> + } + + %overwrites_output = index.cmp eq, %output_accumulation, %c0 : index + %has_one_token = index.cmp eq, %bounded_token_count, %c1 : index + %use_single_token_overwrite = scalar.andi %overwrites_output, %has_one_token : i1 + scf.if %use_single_token_overwrite { + template.apply<@ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32>(%result, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias) : (vector<8xf32>, index, index, index, buffer) + } else { + template.apply<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.body_i8_storage> device @ggml_matmul_vector_q6_f32_body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c17536_bytes = index.constant 17536 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_base = index.mul %output_block, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_base = index.scale %row_group_block, %c17536_bytes : index, offset -> offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %code_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<8x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %block_result = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>) -> (vector<8xf32>) { + %quant_group_next = index.add %quant_group, %c1 : index + %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { + %codes0 = vector.load %code_view[%quant_group, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes1 = vector.load %code_view[%quant_group_next, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes0_f32 = vector.sitofp %codes0 : vector<16xi8> to vector<16xf32> + %codes1_f32 = vector.sitofp %codes1 : vector<16xi8> to vector<16xf32> + %scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + scf.yield %values0, %values1 : vector<16xf16>, vector<16xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1 : vector<16xf16>, vector<16xf16> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs_fragment = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs_fragment = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs_fragment, %rhs_fragment, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next : vector<8xf32> + } + scf.yield %block_result : vector<8xf32> + } + + template.apply<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.return +} + + +template.def<@ggml.matmul.vector_q6_f32.body_selected_outputs> device @ggml_matmul_vector_q6_f32_body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %weight_offset_index = config.get @ggml.matmul.vector.weight_offset : index + %weight_offset = index.cast %weight_offset_index : index to offset + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_selected_count = index.assume %selected_count0 [range(%selected_count0, 1, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %result_stage_bytes = index.constant 2048 : offset + %zero_accumulator = vector.constant 0.0 : vector<16xf16> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %selected_outputs_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %selected_outputs, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%base] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %selected_outputs_view = buffer.view %selected_outputs_noalias[%base] : buffer -> view<64xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%base] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%base] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%base] : buffer -> view<64x16xf16, %activation_fragment_layout> + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %valid_selected_output = index.cmp ult, %weight_local_row, %bounded_selected_count : index + %channel0 = scf.if %valid_selected_output -> (index) { + %selected_output_i32 = view.load %selected_outputs_view[%weight_local_row] : view<64xi32> -> i32 + %selected_output_index0 = index.cast %selected_output_i32 : i32 to index + %selected_output_index = index.assume %selected_output_index0 [range(%selected_output_index0, 0, 262143)] : index + scf.yield %selected_output_index : index + } else { + scf.yield %c0 : index + } + %channel = index.assume %channel0 [range(%channel0, 0, 262143)] : index + %channel_in_range = index.cmp ult, %channel, %bounded_output_size : index + %valid_channel = scalar.andi %valid_selected_output, %channel_in_range : i1 + %selected_output_block = index.div %channel, %c64 : index + %selected_weight_local_row = index.rem %channel, %c64 : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<16xf16>) -> (vector<16xf16>) { + %row_group_block0 = index.mul %selected_output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_relative = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %row_group_byte_base = index.add %weight_offset, %row_group_byte_relative : offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%selected_weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %zero_f32 = scalar.constant 0.0 : f32 + scf.yield %zero_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%selected_weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<16xf16>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %half_next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<16xf16> + } + template.apply<@ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %c64, %c0, %subgroup_channel_add, %result_stage, %output_noalias, %output_accumulation) : (vector<16xf16>, index, index, index, index, index, index, buffer, buffer, index) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom new file mode 100644 index 000000000000..0f87ac2802b4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom @@ -0,0 +1,126 @@ +// Vector-shaped GGML matmul reduction for batch-size-1 F32 outputs. +template.decl @ggml.matmul.vector_publish.next_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) +template.decl @ggml.matmul.vector_activation.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.unary_f32.apply(%arg0: index, %arg1: f32) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%arg0: index, %arg1: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 1, 524288)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.matmul.vector.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @ggml_matmul_vector_f32_f32_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %output_noalias = buffer.assume.noalias %output : buffer + %partition = kernel.workgroup.id : index + %element_bytes = index.constant 4 : offset + %plane_elements = index.mul %token_count, %output_size : index + %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset + %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset + %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %activated = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %published) : (index, vector<4xf32>) -> (vector<4xf32>) + vector.store.mask %activated, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.apply<@ggml.matmul.vector_publish.next_vector4>(%publish_word, %publish_format, %token_count0, %output_size0, %token0, %channel, %activated, %next_output) : (i1, index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @ggml_matmul_vector_f32_f32_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_pair { + %published0 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value0 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value0 : f32 + } + %activated0 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published0) : (index, f32) -> (f32) + view.store %activated0, %output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %published1 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %row1] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value1 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value1 : f32 + } + %activated1 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published1) : (index, f32) -> (f32) + view.store %activated1, %output_view[%token, %row1] : f32, view<[%token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_matmul_vector_f32_f32_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom new file mode 100644 index 000000000000..625812e33558 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom @@ -0,0 +1,177 @@ +template.decl @ggml.matmul.vector_activation.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul.vector_activation.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul.vector_activation.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) + +template.decl @ggml.matmul.vector_publish.next_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) +template.decl @ggml.matmul.vector_activation.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +amdgpu.target @ggml_matmul_vector_postops_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.matmul.vector_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.matmul.vector_postops.output_size : %value: index where [range(%value, 1, 524288)] + +config.decl @ggml.matmul.vector_postops.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.matmul.vector_postops.apply_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector_postops.apply_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @ggml_mul_mat_vector_postops_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %apply_bias0 = config.get @ggml.matmul.vector_postops.apply_bias : index + %apply_residual0 = config.get @ggml.matmul.vector_postops.apply_residual : index + %c1 = index.constant 1 : index + %apply_bias = index.cmp eq, %apply_bias0, %c1 : index + %apply_residual = index.cmp eq, %apply_residual0, %c1 : index + scf.if %publish_word { + %with_bias = scf.if %apply_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul.vector_activation.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %apply_residual -> (vector<4xf32>) { + %residual = template.apply<@ggml.matmul.vector_activation.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %with_bias, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %residual : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul.vector_activation.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.matmul.vector_publish.next_vector4>(%publish_word, %publish_format, %token_count0, %output_size0, %token0, %channel, %published, %next_output) : (i1, index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @ggml_mul_mat_vector_postops_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %apply_bias0 = config.get @ggml.matmul.vector_postops.apply_bias : index + %apply_residual0 = config.get @ggml.matmul.vector_postops.apply_residual : index + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %apply_bias = index.cmp eq, %apply_bias0, %c1 : index + %apply_residual = index.cmp eq, %apply_residual0, %c1 : index + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %bias_noalias, %residual_input_noalias, %residual_output_noalias = buffer.assume.noalias %bias, %residual_input, %residual_output : buffer, buffer, buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %residual_output_view = buffer.view %residual_output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + scf.if %publish_pair { + %biased0 = scf.if %apply_bias -> (f32) { + %bias_value = view.load %bias_view[%channel] : view<[%output_size]xf32> -> f32 + %sum = scalar.addf %value0, %bias_value : f32 + scf.yield %sum : f32 + } else { + scf.yield %value0 : f32 + } + %published0 = scf.if %apply_residual -> (f32) { + %residual = view.load %residual_input_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %biased0, %residual : f32 + scf.yield %sum : f32 + } else { + scf.yield %biased0 : f32 + } + view.store %published0, %residual_output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %biased1 = scf.if %apply_bias -> (f32) { + %bias_value = view.load %bias_view[%row1] : view<[%output_size]xf32> -> f32 + %sum = scalar.addf %value1, %bias_value : f32 + scf.yield %sum : f32 + } else { + scf.yield %value1 : f32 + } + %published1 = scf.if %apply_residual -> (f32) { + %residual = view.load %residual_input_view[%token, %row1] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %biased1, %residual : f32 + scf.yield %sum : f32 + } else { + scf.yield %biased1 : f32 + } + view.store %published1, %residual_output_view[%token, %row1] : f32, view<[%token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_mul_mat_vector_postops_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_bias_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %bias, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_residual_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_bias_residual_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom new file mode 100644 index 000000000000..870419b2b38b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom @@ -0,0 +1,39 @@ +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.def @ggml.matmul.vector.packed_layout = 1 : index + +config.decl @ggml.matmul.vector.q6_storage : %value: index where [range(%value, 0, 1)] + +template.decl @ggml.matmul.vector_q6_f32.body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.def export("ggml_mul_mat_vector_q6_f32_f32") @ggml_mul_mat_vector_q6_f32_f32(%token_count: index) { + %output_size = config.get @ggml.matmul.vector.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_blocks = index.div %padded_output_size, %c64 : index + kernel.launch.config workgroups(%output_blocks, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %scale_row_layout = config.get @ggml.matmul.vector.packed_layout : index + %q6_storage = config.get @ggml.matmul.vector.q6_storage : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %output_block = kernel.workgroup.id : index + %c1 = index.constant 1 : index + %uses_i8_storage = index.cmp eq, %q6_storage, %c1 : index + scf.if %uses_i8_storage { + template.apply<@ggml.matmul.vector_q6_f32.body_i8_storage>(%bounded_token_count, %input_size, %output_size, %output_block, %input, %weight, %output) : (index, index, index, index, buffer, buffer, buffer) + } else { + template.apply<@ggml.matmul.vector_q6_f32.body_packed_storage>(%bounded_token_count, %input_size, %output_size, %output_block, %scale_row_layout, %input, %weight, %output) : (index, index, index, index, index, buffer, buffer, buffer) + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom new file mode 100644 index 000000000000..d8ca1ee69144 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom @@ -0,0 +1,30 @@ +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.def @ggml.matmul.vector.packed_layout = 1 : index + +config.decl @ggml.matmul.vector.weight_offset : %value: index where [range(%value, 0, 2147483647), mul(%value, 256)] + +template.decl @ggml.matmul.vector_q6_f32.body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) + +kernel.def export("ggml_mul_mat_vector_q6_selected_outputs_f32_f32") @ggml_mul_mat_vector_q6_selected_outputs_f32_f32(%token_count: index, %selected_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %selected_count: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %selected_values: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size0 = config.get @ggml.matmul.vector.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %scale_row_layout = config.get @ggml.matmul.vector.packed_layout : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_selected_count = index.assume %selected_count [range(%selected_count, 1, 64)] : index + template.apply<@ggml.matmul.vector_q6_f32.body_selected_outputs>(%bounded_token_count, %input_size, %output_size, %scale_row_layout, %bounded_selected_count, %input, %weight, %selected_outputs, %selected_values) : (index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32>(%bounded_selected_count, %output_size, %selected_outputs, %selected_values, %output) : (index, index, buffer, buffer, buffer) + kernel.return +} diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 51162048f255..005660caaa27 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -1068,7 +1068,7 @@ static void run_rmsnorm_support_checks() { ggml_build_forward_expand(q8_graph, q8_output); require_kernel_subsequence(scheduled_kernel_sequence(q8_graph), { "loom_libs:ggml_rmsnorm_binary_f32", - "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma" }); + "loom_libs:ggml_mul_mat_vector_q6_f32_f32" }); ggml_tensor * wrong_type_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 1); ggml_tensor * wrong_type_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 256); @@ -3039,9 +3039,9 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); + 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), - cache_type != GGML_TYPE_COUNT ? 1.0e-3f : 1.0f, - cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); + abs_tolerance, cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -3054,7 +3054,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, - int64_t output_size) { + int64_t output_size, + int64_t input_size = kQwenHiddenSize) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -3067,7 +3068,6 @@ static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t input_size = kQwenHiddenSize; ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); ggml_tensor * cpu_matmul = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); @@ -3321,6 +3321,114 @@ static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_t ggml_backend_free(hrx_backend); } +static void run_vector_q8_publish_cpu_reference_case() { + constexpr int64_t input_size = 256; + constexpr int64_t intermediate_size = 256; + constexpr int64_t output_size = 128; + constexpr int64_t token_count = 1; + + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 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_first_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, intermediate_size); + ggml_tensor * cpu_second_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_hidden = ggml_mul_mat(cpu_ctx, cpu_first_weight, cpu_input); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_second_weight, cpu_hidden); + ggml_tensor * hrx_first_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, intermediate_size); + ggml_tensor * hrx_second_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_hidden = ggml_mul_mat(hrx_ctx, hrx_first_weight, hrx_input); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_second_weight, hrx_hidden); + 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); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence( + scheduled_kernel_sequence(hrx_graph), + { "loom_libs:ggml_mul_mat_vector_f32_f32", "qwen3_moe:ggml_quantize_q8_1_x4_f32", + "loom_libs:ggml_mul_mat_vector_f32_f32" }); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*hrx_graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const ggml::hrx::Value * hidden_value = imported.graph.values().find_tensor(hrx_hidden); + REQUIRE(hidden_value != nullptr); + const size_t q8_byte_count = static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, intermediate_size); + const ggml::hrx::CommandPlanAlternateValue * q8_alternate = + plan.metadata.find_alternate_value(hidden_value->id, GGML_TYPE_Q8_1, q8_byte_count); + REQUIRE(q8_alternate != nullptr); + + const ggml::hrx::Dispatch * second_vector = nullptr; + int vector_count = 0; + for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_vector_f32_f32") { + ++vector_count; + if (vector_count == 2) { + second_vector = &dispatch; + break; + } + } + } + REQUIRE(second_vector != nullptr); + REQUIRE(second_vector->bindings.size() == 4); + REQUIRE(second_vector->bindings[0].value == q8_alternate->alternate_value); + REQUIRE(second_vector->kernel.compile_parameters.at("ggml.matmul.vector.activation_format") == + std::to_string(GGML_TYPE_Q8_1)); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + 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()); + 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); + 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, 5.0e-2f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + static void run_dense_matmul_postops_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, @@ -4469,15 +4577,61 @@ static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, b ggml_backend_free(backend); } +static void run_vector_output_capacity_experiment() { + constexpr int64_t kSmallVectorInputSize = 256; + constexpr int64_t kAboveOldOutputSize = 262208; + constexpr int64_t kLiftedOutputSize = 524288; + + run_dense_matmul_cpu_reference_case( + GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, kSmallVectorInputSize); + run_dense_matmul_unary_cpu_reference_case( + GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, kSmallVectorInputSize); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_f32_f32", 1, kAboveOldOutputSize, true, false, + false, kSmallVectorInputSize); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kLiftedOutputSize, kSmallVectorInputSize); +} + +static void run_vector_postops_cpu_reference_case() { + constexpr int64_t kVectorInputSize = 256; + constexpr int64_t kVectorOutputSize = 256; + + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_f32_f32", 1, kVectorOutputSize, true, false, + false, kVectorInputSize); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, kVectorOutputSize, true, true, + false, kVectorInputSize); +} + int main() { - run_rmsnorm_support_checks(); - run_alternate_value_alias_lookup_checks(); + const bool run_vector_cap_experiment = std::getenv("GGML_HRX_RUN_VECTOR_CAP_EXPERIMENT") != nullptr; + const bool run_vector_publish_test = std::getenv("GGML_HRX_RUN_VECTOR_PUBLISH_TEST") != nullptr; + const bool run_vector_postops_test = std::getenv("GGML_HRX_RUN_VECTOR_POSTOPS_TEST") != nullptr; + if (!run_vector_cap_experiment && !run_vector_publish_test && !run_vector_postops_test) { + run_rmsnorm_support_checks(); + run_alternate_value_alias_lookup_checks(); + } if (ggml_backend_hrx_get_device_count() == 0) { std::fprintf(stderr, "test skipped: no HRX devices available\n"); return 0; } + if (run_vector_cap_experiment) { + run_vector_output_capacity_experiment(); + return 0; + } + if (run_vector_publish_test) { + run_vector_q8_publish_cpu_reference_case(); + return 0; + } + if (run_vector_postops_test) { + run_vector_postops_cpu_reference_case(); + return 0; + } + run_add_f32_cpu_reference_case(); for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { run_ssm_conv_prefill_recurrent_case(type, false, false); @@ -4514,7 +4668,7 @@ int main() { for (int64_t outputs : { 1, 4, 47, 48, 63 }) { run_dense_matmul_cpu_reference_case( GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" : + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", tokens, outputs); } @@ -4527,42 +4681,42 @@ int main() { GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { for (ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { - run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", + run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_decode_f32_f32", 1, 1024, 2048, false, false, cache_type); } } run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256, 640); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128, 544); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 256, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 640); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 1024, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 1024, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 1024, 640); run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_wmma"); - run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1); + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 256); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); @@ -4584,6 +4738,7 @@ int main() { GGML_GLU_OP_SWIGLU, true); run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + run_vector_q8_publish_cpu_reference_case(); for (ggml_glu_op op : { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { run_dense_matmul_swiglu_cpu_reference_case( @@ -4659,12 +4814,16 @@ int main() { GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, false, true, true, 16640); for (const int64_t tokens : { 1, 3, 5 }) { run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", tokens, 256, 5120, true); + GGML_TYPE_Q4_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, 256, 5120, true); run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", tokens, 256, 5120, true, true); + GGML_TYPE_Q6_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, 256, 5120, true, true); } run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 256, 5120, true); + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 256, 5120, true); run_rmsnorm_cpu_reference_case(); run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); run_strided_rmsnorm_cpu_reference_case(2); From 73afe1eca5c711dc7079f3ab7e0e3fb48669077e Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 10:39:21 -0700 Subject: [PATCH 2/8] Refactor HRX tiled matmul dispatch Move tiled matmul kernels into composable matmul motifs and ops, route covered mul_mat and gated cases through the new names, and move LLM attention matmul wrappers under ops/llm with tiled, vector, and explicit legacy variants. Assisted-by: Codex --- .../common/dispatch-gated-mul-mat.cpp | 218 +++- .../common/dispatch-mul-mat.cpp | 332 +++++- .../llm/dispatch-attention-qkv.cpp | 84 +- .../kernels/loom-libs/manifest.json | 942 +++++++++++++++++- ...iled_input_f32_dequant_accumulate_f32.loom | 232 +++++ ...pair_input_f32_dequant_accumulate_f32.loom | 232 +++++ .../matmul/tiled_pair_postprocess_f32.loom | 84 ++ .../motifs/matmul/tiled_postprocess_f32.loom | 31 + .../motifs/matmul/tiled_publish.loom | 42 + ..._matmul_rope_set_rows_legacy_f32_f32.loom} | 4 +- ..._k_matmul_rope_set_rows_tiled_f32_f32.loom | 72 ++ ...k_matmul_rope_set_rows_vector_f32_f32.loom | 126 +++ ...tention_q_matmul_rope_legacy_f32_f32.loom} | 4 +- ...attention_q_matmul_rope_tiled_f32_f32.loom | 66 ++ ...ttention_q_matmul_rope_vector_f32_f32.loom | 104 ++ ...ion_v_matmul_set_rows_legacy_f32_f32.loom} | 4 +- ...ntion_v_matmul_set_rows_tiled_f32_f32.loom | 63 ++ ...tion_v_matmul_set_rows_vector_f32_f32.loom | 108 ++ ...k_matmul_rope_set_rows_decode_f32_f32.loom | 110 -- ...ttention_q_matmul_rope_decode_f32_f32.loom | 78 -- ...tion_v_matmul_set_rows_decode_f32_f32.loom | 94 -- .../matmul/tiled_input_f32_publish_f32.loom | 263 +++++ ...led_pair_input_f32_binary_publish_f32.loom | 176 ++++ tests/hrx-backend-test.cpp | 446 +++++++-- tests/test-hrx-ops.cpp | 27 +- 25 files changed, 3457 insertions(+), 485 deletions(-) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom => llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom} (96%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{llm_attention_q_matmul_rope_f32_f32_wmma.loom => llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom} (95%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{llm_attention_v_matmul_set_rows_f32_f32_wmma.loom => llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom} (95%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom 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 f6461d422069..b09f45d88d40 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 @@ -17,6 +17,14 @@ static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryBiasPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryAddPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_decode_wave64"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32LowTokenDotKernel = @@ -71,6 +79,24 @@ struct MulMatSwiGLUMatch { bool decode_matched() const { return topology_matched() && token_count == 1; } }; +struct MulMatSwiGLUPostOpsMatch { + MulMatSwiGLUMatch root; + const Value * bias = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const GraphNode * bias_add_node = nullptr; + const GraphNode * residual_add_node = nullptr; + std::vector add_nodes; + KernelCatalogRef kernel = {}; + bool has_bias = false; + bool has_residual = false; + + bool matched() const { + return root.matched() && residual_output != nullptr && kernel.id != kUncatalogedKernelId && + (has_bias || has_residual); + } +}; + struct MulMatSwiGLUProjectionMatch { MulMatSwiGLUMatch gate_up; const GraphNode * projection_node = nullptr; @@ -129,6 +155,19 @@ static bool checked_mul_size(size_t lhs, size_t rhs, size_t & result) { return true; } +static bool is_bias_shape(const Value & value, int64_t output_size) { + if (value.kind != ValueKind::External || value.type != GGML_TYPE_F32 || !value.contiguous || + value.ne[0] != output_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] != 1) { + return false; + } + } + return true; +} + static size_t symmetric_i4_weight_byte_count(int64_t input_size, int64_t output_size) { return static_cast(output_size) * static_cast(input_size / 256) * size_t{ 144 }; } @@ -247,7 +286,7 @@ static MulMatSwiGLUMatch match_mul_mat_swiglu(const DispatchMatchContext & conte static PackedMulMatGluMatch match_packed_mul_mat_glu(const DispatchMatchContext & context) { PackedMulMatGluMatch match; const CommonMulMatMatch root = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSwiGLUF32F32WmmaKernel, false); + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledPairF32BinaryPublishF32Kernel, false); if (!root.matched() || !context.graph.has_index() || root.weight->alias_source.value >= 0) { return {}; } @@ -314,7 +353,7 @@ static bool match_packed_mul_mat_glu_dispatch(const DispatchMatchContext & conte } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUF32F32WmmaKernel); + dispatch.kernel = make_kernel_specialization(kMulMatTiledPairF32BinaryPublishF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_to_config_value(match.token_count)); @@ -671,11 +710,166 @@ static bool has_qualified_swiglu_q8_consumer(const Graph & graph, const MulMatSw return false; } +static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLUMatch & match) { + if (!match.matched()) { + 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 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; +} + +static MulMatSwiGLUPostOpsMatch match_mul_mat_swiglu_postops(const DispatchMatchContext & context) { + MulMatSwiGLUPostOpsMatch match; + match.root = match_mul_mat_swiglu(context); + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match.root) || !context.graph.has_index()) { + return {}; + } + + const Value * current = match.root.output; + for (int add_index = 0; add_index < 2; ++add_index) { + const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); + if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { + break; + } + + const bool current_is_lhs = add_node->inputs[0] == current->id; + const bool current_is_rhs = add_node->inputs[1] == current->id; + if (!current_is_lhs && !current_is_rhs) { + return {}; + } + + const ValueId other_id = current_is_lhs ? add_node->inputs[1] : add_node->inputs[0]; + const Value * other = common_graph_value(context.graph, other_id); + const Value * output = common_graph_value(context.graph, add_node->output); + if (other == nullptr || output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_same_shape(*match.root.output, *output)) { + return {}; + } + + if (!match.has_bias && is_bias_shape(*other, match.root.output_size)) { + match.bias = other; + match.bias_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_bias = true; + current = output; + continue; + } + + if (!match.has_residual && other->type == GGML_TYPE_F32 && other->contiguous && + common_same_shape(*match.root.output, *other)) { + match.residual_input = other; + match.residual_output = output; + match.residual_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_residual = true; + current = output; + continue; + } + + break; + } + + if (!match.has_bias && !match.has_residual) { + return {}; + } + + match.residual_output = current; + if (match.has_bias && match.has_residual) { + match.kernel = kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel; + } else if (match.has_residual) { + match.kernel = kMulMatTiledPairF32BinaryAddPublishF32Kernel; + } else if (match.has_bias) { + match.kernel = kMulMatTiledPairF32BinaryBiasPublishF32Kernel; + } + return match; +} + +static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, + const MulMatSwiGLUPostOpsMatch * postops, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch & match = postops != nullptr ? postops->root : match_mul_mat_swiglu(context); + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(postops != nullptr ? postops->kernel : + kMulMatTiledPairF32BinaryPublishF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.gate_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.up_format))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op", + common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + if (postops != nullptr && postops->has_bias) { + dispatch.bindings.push_back({ postops->bias->id, 0, postops->bias->byte_count }); + } + if (postops != nullptr && postops->has_residual) { + dispatch.bindings.push_back({ postops->residual_input->id, 0, postops->residual_input->byte_count }); + } + const Value * output = postops != nullptr ? postops->residual_output : match.output; + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.up_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.glu_node, + dispatch_match.covered_nodes)) { + return false; + } + if (postops != nullptr) { + for (const GraphNode * add_node : postops->add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_tiled_pair_mul_mat_swiglu_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUPostOpsMatch match = match_mul_mat_swiglu_postops(context); + if (!match.matched()) { + return false; + } + return build_tiled_pair_mul_mat_swiglu_dispatch(context, &match, dispatch_match); +} + +static bool match_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + return build_tiled_pair_mul_mat_swiglu_dispatch(context, nullptr, dispatch_match); +} + static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); if (!match.matched()) { return false; } + if (common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + 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 && @@ -822,7 +1016,23 @@ void register_gated_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_mul_mat_swiglu_q5_projection_dispatch, }); registry.add({ - "common.mul_mat_swiglu.f32_f32_wmma", + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_tiled_pair_mul_mat_swiglu_postops_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.tiled_pair_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_tiled_pair_mul_mat_swiglu_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.legacy_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 290, @@ -830,7 +1040,7 @@ void register_gated_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_mul_mat_swiglu_dispatch, }); registry.add({ - "common.packed_mul_mat_glu.f32_f32_wmma", + "common.packed_mul_mat_glu.tiled_pair_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 285, 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 d7ffd80bd75f..5c55a3b4a357 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 @@ -15,6 +15,10 @@ namespace { static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledF32F16AlternateKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate"); static constexpr KernelCatalogRef kMulMatF32F32NarrowSplitK4Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_narrow_split_k4"); static constexpr KernelCatalogRef kMulMatF32F32DecodeWave64Kernel = @@ -37,6 +41,12 @@ static constexpr KernelCatalogRef kMulMatVectorBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_residual_f32_f32"); static constexpr KernelCatalogRef kQuantizeQ8_1X4F32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); +static constexpr KernelCatalogRef kMulMatTiledBiasF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledBiasAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"); static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K64PlaneKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i4_k64_plane"); static constexpr KernelCatalogRef kMulMatSymmetricI4WmmaKernel = @@ -135,6 +145,35 @@ struct VectorPublishPlan { size_t q8_byte_count = 0; }; +static bool common_mul_mat_uses_narrow_split_k_route(const CommonMulMatMatch & match) { + return match.weight_format == CommonMulMatWeightFormat::Q4K && !match.has_fused_unary && + match.input_size >= 4096 && match.input_size % 1024 == 0 && match.output_size <= 64 && + match.output_size % 4 == 0 && match.token_count >= 128 && match.token_count % 32 == 0 && + common_ceil_div(match.token_count, 32) < 32; +} + +static bool common_mul_mat_is_direct_f32_tiled_route(const CommonMulMatMatch & match) { + return match.matched() && match.token_count > 5 && !common_mul_mat_uses_narrow_split_k_route(match); +} + +static bool common_mul_mat_is_skinny_route(const CommonMulMatMatch & match) { + return match.matched() && match.token_count <= 5; +} + +static bool common_mul_mat_postops_is_direct_f32_tiled_route(const MulMatPostOpsMatch & match) { + return match.matched() && match.layout_node == nullptr && !match.requires_q8_activation && match.token_count > 5; +} + +static bool common_mul_mat_postops_is_skinny_route(const MulMatPostOpsMatch & match) { + return match.matched() && match.token_count <= 5; +} + +static bool common_mul_mat_uses_factored_tiled_kernel(KernelCatalogRef kernel) { + return kernel.id == kMulMatTiledF32F32Kernel.id || kernel.id == kMulMatTiledF32F16AlternateKernel.id || + kernel.id == kMulMatTiledBiasF32F32Kernel.id || kernel.id == kMulMatTiledAddF32F32Kernel.id || + kernel.id == kMulMatTiledBiasAddF32F32Kernel.id; +} + static bool is_bias_shape(const Value & value, int64_t output_size) { if (value.kind != ValueKind::External || value.type != GGML_TYPE_F32 || !value.contiguous || value.ne[0] != output_size) { @@ -148,6 +187,54 @@ static bool is_bias_shape(const Value & value, int64_t output_size) { return true; } +static bool common_mul_mat_has_f16_alternate_consumer(const Graph & graph, const CommonMulMatMatch & match) { + if (!graph.has_index() || match.output == nullptr || match.output->type != GGML_TYPE_F32 || + !match.output->contiguous || match.token_count < 128) { + return false; + } + + for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { + const CommonMulMatMatch downstream = + common_match_mul_mat_any_format(graph, consumer, kMulMatQ4KF16WmmaPrefillWave32Kernel, false); + if (!downstream.matched() || downstream.input->id != match.output->id || + downstream.input_size != match.output_size || downstream.token_count != match.token_count || + downstream.output_size % 64 != 0 || downstream.input_size % 256 != 0) { + continue; + } + if (downstream.weight->type == GGML_TYPE_Q6_K && downstream.weight->alias_source.value < 0 && + downstream.token_count % 128 == 0) { + return true; + } + if (downstream.weight->type == GGML_TYPE_Q4_K && downstream.weight->alias_source.value < 0 && + downstream.token_count % 256 == 0 && downstream.output_size >= downstream.input_size / 4) { + return true; + } + } + return false; +} + +static bool common_mul_mat_has_q8_alternate_consumer(const Graph & graph, const CommonMulMatMatch & match) { + if (!graph.has_index() || match.output == nullptr || match.output->type != GGML_TYPE_F32 || + !match.output->contiguous || match.token_count < 256 || match.token_count > 2048 || + match.token_count % 256 != 0 || match.output_size % 128 != 0) { + return false; + } + + for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { + const CommonMulMatMatch downstream = + common_match_mul_mat_any_format(graph, consumer, kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel, false); + if (!downstream.matched() || downstream.input->id != match.output->id || + downstream.input_size != match.output_size || downstream.token_count != match.token_count || + downstream.output_size % 64 != 0 || downstream.input_size % 256 != 0) { + continue; + } + if (downstream.weight->type == GGML_TYPE_Q5_K || downstream.weight->type == GGML_TYPE_IQ4_XS) { + return true; + } + } + return false; +} + static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) { if (!graph.has_index()) { return false; @@ -912,12 +999,15 @@ static bool match_q4_k_q8_1_x4_prefill_dispatch(const DispatchMatchContext & con } static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, - bool allow_batch1_bias_residual) { + KernelCatalogRef root_kernel, + KernelCatalogRef bias_kernel, + KernelCatalogRef add_kernel, + KernelCatalogRef bias_add_kernel, + bool allow_batch1_bias_residual = false) { MulMatPostOpsMatch match; - CommonMulMatMatch root = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + CommonMulMatMatch root = common_match_mul_mat_any_format(context.graph, context.root_node, root_kernel, false); if (!root.matched()) { - root = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, true); + root = common_match_mul_mat_any_format(context.graph, context.root_node, root_kernel, true); } if (!root.matched() || !context.graph.has_index()) { return match; @@ -991,11 +1081,11 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con match.residual_output = current; if (match.has_bias && match.has_residual) { - match.kernel = kMulMatBiasAddF32F32WmmaKernel; + match.kernel = bias_add_kernel; } else if (match.has_residual) { - match.kernel = kMulMatAddF32F32WmmaKernel; + match.kernel = add_kernel; } else if (match.has_bias) { - match.kernel = kMulMatBiasF32F32WmmaKernel; + match.kernel = bias_kernel; } else { return {}; } @@ -1012,7 +1102,19 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con } static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context) { - return match_mul_mat_postops(context, false); + return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, + kMulMatBiasF32F32WmmaKernel, + kMulMatAddF32F32WmmaKernel, + kMulMatBiasAddF32F32WmmaKernel); +} + +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, + bool allow_batch1_bias_residual) { + return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, + kMulMatVectorBiasF32F32Kernel, + kMulMatVectorAddF32F32Kernel, + kMulMatVectorBiasAddF32F32Kernel, + allow_batch1_bias_residual); } static bool try_match_fused_unary(const DispatchMatchContext & context, CommonMulMatMatch & match) { @@ -1058,13 +1160,16 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, const CommonMulMatMatch & match, DispatchMatch & dispatch_match, CommonQ8ActivationPolicy q8_policy = CommonQ8ActivationPolicy::ExistingAlternateOnly) { - const bool split_k = match.weight_format == CommonMulMatWeightFormat::Q4K && !match.has_fused_unary && - match.input_size >= 4096 && match.input_size % 1024 == 0 && - match.output_size <= 64 && match.output_size % 4 == 0 && - match.token_count >= 128 && match.token_count % 32 == 0 && - common_ceil_div(match.token_count, 32) < 32; + const bool split_k = common_mul_mat_uses_narrow_split_k_route(match); + const bool publish_f16_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && + common_mul_mat_has_f16_alternate_consumer(context.graph, match); + const bool publish_q8_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && + common_mul_mat_has_q8_alternate_consumer(context.graph, match); Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(split_k ? kMulMatF32F32NarrowSplitK4Kernel : match.kernel); + dispatch.kernel = make_kernel_specialization( + split_k ? kMulMatF32F32NarrowSplitK4Kernel : + publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : + match.kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_to_config_value(match.token_count)); @@ -1080,7 +1185,15 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, bool pack_q4 = false; bool pack_q6 = false; DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; - if (can_pack_q4 || can_pack_q6) { + const size_t f16_activation_bytes = static_cast(match.token_count * match.input_size) * sizeof(ggml_fp16_t); + const CommandPlanAlternateValue * f16_activation = + common_mul_mat_uses_factored_tiled_kernel(match.kernel) ? + find_alternate_value(context.graph, context.plan, match.input->id, GGML_TYPE_F16, f16_activation_bytes) : + nullptr; + if (f16_activation != nullptr) { + activation = { f16_activation->alternate_value, 0, f16_activation_bytes }; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); + } else if (can_pack_q4 || can_pack_q6) { if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation, q8_policy)) { pack_q4 = can_pack_q4; @@ -1120,12 +1233,48 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, } else { dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); } + if (publish_f16_alternate) { + const size_t bytes = match.output->byte_count / 2; + const ValueId f16_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + constexpr const char * name = "common.mul_mat.tiled.f16"; + dispatch_match.transients.push_back({ f16_output, name, bytes, 256 }); + Status status; + if (!dispatch_match.metadata.append_alternate_value( + { match.output->id, f16_output, GGML_TYPE_F16, bytes, name }, status)) { + dispatch_match.status.append(status); + return false; + } + dispatch.bindings.push_back({ f16_output, 0, bytes }); + } dispatch_match.covered_nodes.push_back(context.root_index); if (match.has_fused_unary) { dispatch_match.covered_nodes.push_back(match.unary_node_index); } dispatch_match.dispatches.push_back(std::move(dispatch)); + if (publish_q8_alternate) { + const size_t bytes = static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.output_size); + const ValueId q8_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + constexpr const char * name = "common.mul_mat.tiled.q8_1_x4"; + dispatch_match.transients.push_back({ q8_output, name, bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeQ8_1X4F32Kernel); + quantize.kernel.integer_parameters.emplace("token_count", match.token_count); + quantize.kernel.integer_parameters.emplace("input_size", match.output_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + common_to_config_value(match.token_count * match.output_size / 128)); + quantize.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + quantize.bindings.push_back({ q8_output, 0, bytes }); + + Status status; + if (!dispatch_match.metadata.append_alternate_value( + { match.output->id, q8_output, GGML_TYPE_Q8_1, bytes, name }, status)) { + dispatch_match.status.append(status); + return false; + } + dispatch_match.dispatches.push_back(std::move(quantize)); + } return true; } @@ -1587,8 +1736,9 @@ static bool match_q6_vector_final_projection_q8_dispatch(const DispatchMatchCont return build_q6_vector_dispatch(match, context.root_index, dispatch_match); } -static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const MulMatPostOpsMatch match = match_mul_mat_postops(context); +static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, + const MulMatPostOpsMatch & match, + DispatchMatch & dispatch_match) { if (!match.matched()) { return false; } @@ -1670,6 +1820,41 @@ static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, return true; } +static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, + kMulMatBiasF32F32WmmaKernel, + kMulMatAddF32F32WmmaKernel, + kMulMatBiasAddF32F32WmmaKernel); + if (common_mul_mat_postops_is_direct_f32_tiled_route(match) || common_mul_mat_postops_is_skinny_route(match)) { + return false; + } + return build_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static bool match_skinny_mul_mat_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, + kMulMatBiasF32F32WmmaKernel, + kMulMatAddF32F32WmmaKernel, + kMulMatBiasAddF32F32WmmaKernel); + if (!common_mul_mat_postops_is_skinny_route(match)) { + return false; + } + return build_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static bool match_tiled_mul_mat_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatTiledF32F32Kernel, + kMulMatTiledBiasF32F32Kernel, + kMulMatTiledAddF32F32Kernel, + kMulMatTiledBiasAddF32F32Kernel); + if (!common_mul_mat_postops_is_direct_f32_tiled_route(match)) { + return false; + } + return build_mul_mat_postops_dispatch(context, match, dispatch_match); +} + static DecodeMulMatAddMatch match_decode_mul_mat_add(const DispatchMatchContext & context) { DecodeMulMatAddMatch match; CommonMulMatMatch root = @@ -1772,6 +1957,35 @@ static bool match_mul_mat_dispatch(const DispatchMatchContext & context, Dispatc return false; } try_match_fused_unary(context, match); + if (common_mul_mat_is_direct_f32_tiled_route(match) || common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_skinny_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + if (!match.matched()) { + return false; + } + try_match_fused_unary(context, match); + if (!common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); +} + +static bool match_tiled_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); + if (!match.matched()) { + return false; + } + try_match_fused_unary(context, match); + if (!common_mul_mat_is_direct_f32_tiled_route(match)) { + return false; + } return build_mul_mat_dispatch(context, match, dispatch_match); } @@ -1781,6 +1995,30 @@ static bool match_mul_mat_unary_dispatch(const DispatchMatchContext & context, D if (!match.matched() || !try_match_fused_unary(context, match)) { return false; } + if (common_mul_mat_is_direct_f32_tiled_route(match) || common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); +} + +static bool match_tiled_mul_mat_unary_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); + if (!match.matched() || !try_match_fused_unary(context, match) || + !common_mul_mat_is_direct_f32_tiled_route(match)) { + return false; + } return build_mul_mat_dispatch(context, match, dispatch_match); } @@ -1797,7 +2035,7 @@ static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, match.weight_format == CommonMulMatWeightFormat::Q6K) && match.output_size % 64 == 0 && common_is_supported_dense_output_size(match.output_size) && match.weight->alias_source.value < 0) { match.kernel = kMulMatF32F32WmmaKernel; - return build_mul_mat_dispatch(context, match, dispatch_match); + return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); } else { build_decode_mul_mat_dispatch(match, dispatch_match, context.root_index); } @@ -1917,7 +2155,23 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_vector_mul_mat_unary_dispatch, }); registry.add({ - "common.mul_mat_unary.f32_f32_wmma", + "common.mul_mat_unary.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_tiled_mul_mat_unary_dispatch, + }); + registry.add({ + "common.mul_mat_unary.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_skinny_mul_mat_unary_dispatch, + }); + registry.add({ + "common.mul_mat_unary.legacy_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 290, @@ -1933,7 +2187,23 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_vector_mul_mat_dispatch, }); registry.add({ - "common.mul_mat_postops.f32_f32_wmma", + "common.mul_mat_postops.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 181, + DispatchSource::Common, + match_tiled_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat_postops.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 182, + DispatchSource::Common, + match_skinny_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat_postops.legacy_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 180, @@ -1941,7 +2211,23 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_mul_mat_postops_dispatch, }); registry.add({ - "common.mul_mat.f32_f32_wmma", + "common.mul_mat.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 81, + DispatchSource::Common, + match_tiled_mul_mat_dispatch, + }); + registry.add({ + "common.mul_mat.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 82, + DispatchSource::Common, + match_skinny_mul_mat_dispatch, + }); + registry.add({ + "common.mul_mat.legacy_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 80, @@ -1949,7 +2235,7 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_mul_mat_dispatch, }); registry.add({ - "common.mul_mat_add.f32_f32_decode", + "common.mul_mat_add.skinny_f32_f32_decode", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 70, @@ -1957,7 +2243,7 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { match_decode_mul_mat_add_dispatch, }); registry.add({ - "common.mul_mat.f32_f32_decode", + "common.mul_mat.skinny_f32_f32_decode", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 60, 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 ea81bda320f3..5cbc4e0e53f9 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 @@ -15,18 +15,24 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kAttentionQMatMulRopeF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_f32_f32_wmma"); -static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma"); -static constexpr KernelCatalogRef kAttentionVMatMulSetRowsF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_f32_f32_wmma"); -static constexpr KernelCatalogRef kAttentionQMatMulRopeF32F32DecodeKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_decode_f32_f32"); -static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsF32F32DecodeKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_decode_f32_f32"); -static constexpr KernelCatalogRef kAttentionVMatMulSetRowsF32F32DecodeKernel = - GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_decode_f32_f32"); +static constexpr KernelCatalogRef kAttentionQMatMulRopeTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionQMatMulRopeVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionQMatMulRopeLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_legacy_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_legacy_f32_f32"); static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); @@ -455,22 +461,22 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte match.layout_nodes.insert(match.layout_nodes.end(), set_rows_layouts.begin(), set_rows_layouts.end()); match.root = root; match.kernel = is_supported_decode_token_count(root.token_count) ? - kAttentionKMatMulRopeSetRowsF32F32DecodeKernel : - kAttentionKMatMulRopeSetRowsF32F32WmmaKernel; + kAttentionKMatMulRopeSetRowsVectorF32F32Kernel : + kAttentionKMatMulRopeSetRowsTiledF32F32Kernel; if (can_use_packed_lowtoken_attention_matmul(root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionKMatMulRopeSetRowsF32F32WmmaKernel; + match.kernel = kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel; } match.kind = AttentionQkvProjectionKind::Key; return match; } match.root = root; - match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionQMatMulRopeF32F32DecodeKernel : - kAttentionQMatMulRopeF32F32WmmaKernel; + match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionQMatMulRopeVectorF32F32Kernel : + kAttentionQMatMulRopeTiledF32F32Kernel; if (can_use_packed_lowtoken_attention_matmul(root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionQMatMulRopeF32F32WmmaKernel; + match.kernel = kAttentionQMatMulRopeLegacyF32F32Kernel; } match.kind = AttentionQkvProjectionKind::Query; return match; @@ -487,11 +493,11 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte return {}; } match.root = root; - match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionVMatMulSetRowsF32F32DecodeKernel : - kAttentionVMatMulSetRowsF32F32WmmaKernel; + match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionVMatMulSetRowsVectorF32F32Kernel : + kAttentionVMatMulSetRowsTiledF32F32Kernel; if (can_use_packed_lowtoken_attention_matmul(root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionVMatMulSetRowsF32F32WmmaKernel; + match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; } match.kind = AttentionQkvProjectionKind::Value; return match; @@ -535,12 +541,6 @@ static void add_attention_qkv_compile_parameters(Dispatch & dispatch, const Atte dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.input_size", to_config_value(match.root.input_size)); dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.output_size", to_config_value(match.root.output_size)); - if (is_supported_decode_token_count(match.root.token_count)) { - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.token_capacity", - to_config_value(match.root.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.output_capacity", - to_config_value(match.root.output_size)); - } dispatch.kernel.compile_parameters.emplace( "llm.attention_qkv.weight_format", to_config_value(common_mul_mat_format_config_value(match.root.weight_format))); @@ -608,18 +608,22 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D match.root.weight_format = fallback_format; pack_q4 = false; pack_q6 = false; - if (is_supported_decode_token_count(match.root.token_count)) { - switch (match.kind) { - case AttentionQkvProjectionKind::Query: - match.kernel = kAttentionQMatMulRopeF32F32DecodeKernel; - break; - case AttentionQkvProjectionKind::Key: - match.kernel = kAttentionKMatMulRopeSetRowsF32F32DecodeKernel; - break; - case AttentionQkvProjectionKind::Value: - match.kernel = kAttentionVMatMulSetRowsF32F32DecodeKernel; - break; - } + switch (match.kind) { + case AttentionQkvProjectionKind::Query: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionQMatMulRopeVectorF32F32Kernel : + kAttentionQMatMulRopeTiledF32F32Kernel; + break; + case AttentionQkvProjectionKind::Key: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionKMatMulRopeSetRowsVectorF32F32Kernel : + kAttentionKMatMulRopeSetRowsTiledF32F32Kernel; + break; + case AttentionQkvProjectionKind::Value: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionVMatMulSetRowsVectorF32F32Kernel : + kAttentionVMatMulSetRowsTiledF32F32Kernel; + break; } } } @@ -665,7 +669,7 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D void register_llm_attention_qkv_dispatches(DispatchRegistryBuilder & registry) { registry.add({ - "llm.attention_qkv_matmul_postprocess.f32_f32_wmma", + "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, 210, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 60e5d4557cef..7549d57afc2f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -45,22 +45,31 @@ "path": "motifs/matmul/vector_decode_pair.loom" }, { - "path": "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom" + "path": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" }, { - "path": "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom" + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" }, { - "path": "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom" + "path": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" }, { - "path": "ops/llm_attention_q_matmul_rope_decode_f32_f32.loom" + "path": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom" }, { - "path": "ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom" + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom" }, { - "path": "ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom" + "path": "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom" }, { "path": "ops/mul_mat_id_f32_f32_wmma.loom" @@ -250,6 +259,27 @@ }, { "path": "motifs/f32_f16.loom" + }, + { + "path": "ops/matmul/tiled_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + }, + { + "path": "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul/tiled_postprocess_f32.loom" + }, + { + "path": "motifs/matmul/tiled_publish.loom" + }, + { + "path": "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul/tiled_pair_postprocess_f32.loom" } ], "exports": [ @@ -1854,10 +1884,10 @@ ] }, { - "name": "llm_attention_q_matmul_rope_f32_f32_wmma", + "name": "llm_attention_q_matmul_rope_tiled_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_q_matmul_rope_f32_f32_wmma", - "source": "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom", + "symbol": "llm_attention_q_matmul_rope_tiled_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -1890,14 +1920,13 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom" + "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/rope_f32.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -1908,11 +1937,10 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/rope_f32.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -1923,10 +1951,10 @@ ] }, { - "name": "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", + "name": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "source": "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom", + "symbol": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -1961,14 +1989,13 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom" + "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/rope_f32.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -1979,11 +2006,10 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/rope_f32.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -1994,10 +2020,10 @@ ] }, { - "name": "llm_attention_v_matmul_set_rows_f32_f32_wmma", + "name": "llm_attention_v_matmul_set_rows_tiled_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_v_matmul_set_rows_f32_f32_wmma", - "source": "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom", + "symbol": "llm_attention_v_matmul_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2026,13 +2052,12 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom" + "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -2043,10 +2068,9 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -2057,10 +2081,10 @@ ] }, { - "name": "llm_attention_q_matmul_rope_decode_f32_f32", + "name": "llm_attention_q_matmul_rope_vector_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_q_matmul_rope_decode_f32_f32", - "source": "ops/llm_attention_q_matmul_rope_decode_f32_f32.loom", + "symbol": "llm_attention_q_matmul_rope_vector_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2093,11 +2117,14 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_q_matmul_rope_decode_f32_f32.loom" + "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom" ], "library_sources": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", "motifs/rope_f32.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2109,8 +2136,11 @@ ] }, "compile_dependencies": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", "motifs/rope_f32.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2122,10 +2152,10 @@ ] }, { - "name": "llm_attention_k_matmul_rope_set_rows_decode_f32_f32", + "name": "llm_attention_k_matmul_rope_set_rows_vector_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_k_matmul_rope_set_rows_decode_f32_f32", - "source": "ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom", + "symbol": "llm_attention_k_matmul_rope_set_rows_vector_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2160,11 +2190,14 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom" + "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom" ], "library_sources": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", "motifs/rope_f32.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2176,8 +2209,11 @@ ] }, "compile_dependencies": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", "motifs/rope_f32.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2189,10 +2225,10 @@ ] }, { - "name": "llm_attention_v_matmul_set_rows_decode_f32_f32", + "name": "llm_attention_v_matmul_set_rows_vector_f32_f32", "family": "loom_libs", - "symbol": "llm_attention_v_matmul_set_rows_decode_f32_f32", - "source": "ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom", + "symbol": "llm_attention_v_matmul_set_rows_vector_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2221,10 +2257,13 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom" + "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom" ], "library_sources": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2236,8 +2275,214 @@ ] }, "compile_dependencies": [ - "ops/mul_mat_f32_f32_decode.loom", + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_q_matmul_rope_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_q_matmul_rope_legacy_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "indices", + "theta", + "freq_factors", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_v_matmul_set_rows_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_v_matmul_set_rows_legacy_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", @@ -5959,6 +6204,607 @@ "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_bias_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_bias_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_residual_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "bias", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] } ], "link_modules": [], diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..57a9543eb066 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,232 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.mul_mat.activation_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Shared contiguous 2D F32 activation tiled matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %bounded_token_count, %c31 : index + %token_tile_count = index.div %padded_token_count, %c32 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %bounded_input_size) : (index, index) -> (offset) + %split_count = kernel.workgroup.count : index + %partition = kernel.workgroup.id : index + %blocks_per_partition = index.div %quant_block_count, %split_count : index + %quant_block_begin = index.mul %partition, %blocks_per_partition : index + %quant_block_end = index.add %quant_block_begin, %blocks_per_partition : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f16_format = index.constant 1 : index + %uses_f16_activation = index.cmp eq, %activation_format, %f16_format : index + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %input_f32_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %input_f16_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf16> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %converted = scf.if %uses_f16_activation -> (vector<4xf16>) { + %loaded = vector.load.mask %input_f16_view[%bounded_token, %input_k], %mask, %c0_f16x4 : view<[%bounded_token_count]x[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + scf.yield %loaded : vector<4xf16> + } else { + %loaded = vector.load.mask %input_f32_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %narrow = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %narrow : vector<4xf16> + } + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile>(%bounded_token_count, %bounded_output_size, %output_tile_count, %token_tile_count, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + + + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + func.call @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..3dde5e2c456b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,232 @@ +// Shared paired tiled matmul pipeline for contiguous 2D F32 activations. +// The core preserves the legacy generic paired tile schedule and delegates +// pair postprocess and publication to the op wrapper. +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) + +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) + +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) + %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) + %input_noalias, %gate_weight_noalias, %up_weight_noalias = buffer.assume.noalias %input, %gate_weight, %up_weight : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %gate_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %up_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %gate_weight_stage_view = buffer.view %gate_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %up_weight_stage_view = buffer.view %up_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %gate_result_fragment_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %up_result_fragment_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %gate_result_physical_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %up_result_physical_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%gate_acc00 = %gate_block_acc00 : vector<8xf16>, %gate_acc01 = %gate_block_acc01 : vector<8xf16>, %gate_acc10 = %gate_block_acc10 : vector<8xf16>, %gate_acc11 = %gate_block_acc11 : vector<8xf16>, %up_acc00 = %up_block_acc00 : vector<8xf16>, %up_acc01 = %up_block_acc01 : vector<8xf16>, %up_acc10 = %up_block_acc10 : vector<8xf16>, %up_acc11 = %up_block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %gate_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %gate_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %up_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %up_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %gate_weight_values, %gate_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + vector.store %up_weight_values, %up_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %loaded = vector.load.mask %input_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 = scf.for %k_half = [%c0 to %c32 step %c16](%gate_half_acc00 = %gate_acc00 : vector<8xf16>, %gate_half_acc01 = %gate_acc01 : vector<8xf16>, %gate_half_acc10 = %gate_acc10 : vector<8xf16>, %gate_half_acc11 = %gate_acc11 : vector<8xf16>, %up_half_acc00 = %up_acc00 : vector<8xf16>, %up_half_acc01 = %up_acc01 : vector<8xf16>, %up_half_acc10 = %up_acc10 : vector<8xf16>, %up_half_acc11 = %up_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %gate_lhs0 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %gate_lhs1 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs0 = vector.fragment.load %up_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs1 = vector.fragment.load %up_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %gate_half_next00 = vector.mma %gate_lhs0, %rhs0, %gate_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next01 = vector.mma %gate_lhs0, %rhs1, %gate_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next10 = vector.mma %gate_lhs1, %rhs0, %gate_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next11 = vector.mma %gate_lhs1, %rhs1, %gate_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next00 = vector.mma %up_lhs0, %rhs0, %up_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next01 = vector.mma %up_lhs0, %rhs1, %up_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next10 = vector.mma %up_lhs1, %rhs0, %up_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next11 = vector.mma %up_lhs1, %rhs1, %up_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %gate_half_next00, %gate_half_next01, %gate_half_next10, %gate_half_next11, %up_half_next00, %up_half_next01, %up_half_next10, %up_half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %gate_result00, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result00, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result01, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result01, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result10, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result10, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result11, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result11, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + func.return +} + +template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32_core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { + func.call @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op, %epilogue, %gate_weight_format, %up_weight_format, %token_count, %input_size0, %output_size0, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom new file mode 100644 index 000000000000..ed61ece2dc27 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom @@ -0,0 +1,84 @@ +// Pairwise tiled matmul postprocess for gate/up projections. +template.decl @ggml.matmul.tiled_pair_postprocess_f32.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.def<@ggml.matmul.tiled_pair_postprocess_f32.apply_vector4> device @ggml_matmul_tiled_pair_postprocess_f32_apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) { + %c0_5_f32x4 = vector.constant 0.5 : vector<4xf32> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %c2_f32x4 = vector.constant 2.0 : vector<4xf32> + %cn2_f32x4 = vector.constant -2.0 : vector<4xf32> + %gelu_coef = vector.constant 0.044715 : vector<4xf32> + %sqrt_2_over_pi = vector.constant 0.79788456080286544 : vector<4xf32> + + %op_add = index.constant 0 : index + %sum = vector.addf %lhs, %rhs : vector<4xf32> + %is_add = index.cmp eq, %op, %op_add : index + %add_selected = scf.select %is_add, %sum, %lhs : vector<4xf32> + + %op_sub = index.constant 1 : index + %difference = vector.subf %lhs, %rhs : vector<4xf32> + %is_sub = index.cmp eq, %op, %op_sub : index + %sub_selected = scf.select %is_sub, %difference, %add_selected : vector<4xf32> + + %op_mul = index.constant 2 : index + %product = vector.mulf %lhs, %rhs : vector<4xf32> + %is_mul = index.cmp eq, %op, %op_mul : index + %mul_selected = scf.select %is_mul, %product, %sub_selected : vector<4xf32> + + %op_div = index.constant 3 : index + %quotient = vector.divf %lhs, %rhs : vector<4xf32> + %is_div = index.cmp eq, %op, %op_div : index + %div_selected = scf.select %is_div, %quotient, %mul_selected : vector<4xf32> + + %op_swiglu = index.constant 4 : index + %negative_lhs = vector.subf %c0_f32x4, %lhs : vector<4xf32> + %lhs_exp = vector.expf %negative_lhs : vector<4xf32> + %denominator = vector.addf %c1_f32x4, %lhs_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32x4, %denominator : vector<4xf32> + %silu = vector.mulf %lhs, %sigmoid : vector<4xf32> + %swiglu = vector.mulf %silu, %rhs : vector<4xf32> + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %swiglu_selected = scf.select %is_swiglu, %swiglu, %div_selected : vector<4xf32> + + %op_geglu = index.constant 5 : index + %gelu_x2 = vector.mulf %lhs, %lhs : vector<4xf32> + %gelu_poly0 = vector.mulf %gelu_coef, %gelu_x2 : vector<4xf32> + %gelu_poly = vector.addf %c1_f32x4, %gelu_poly0 : vector<4xf32> + %gelu_inner0 = vector.mulf %lhs, %gelu_poly : vector<4xf32> + %gelu_inner = vector.mulf %sqrt_2_over_pi, %gelu_inner0 : vector<4xf32> + %gelu_tanh_scaled = vector.mulf %cn2_f32x4, %gelu_inner : vector<4xf32> + %gelu_tanh_exp = vector.expf %gelu_tanh_scaled : vector<4xf32> + %gelu_tanh_denominator = vector.addf %c1_f32x4, %gelu_tanh_exp : vector<4xf32> + %gelu_tanh_ratio = vector.divf %c2_f32x4, %gelu_tanh_denominator : vector<4xf32> + %gelu_tanh = vector.subf %gelu_tanh_ratio, %c1_f32x4 : vector<4xf32> + %gelu_one_plus = vector.addf %c1_f32x4, %gelu_tanh : vector<4xf32> + %gelu_half_x = vector.mulf %c0_5_f32x4, %lhs : vector<4xf32> + %gelu = vector.mulf %gelu_half_x, %gelu_one_plus : vector<4xf32> + %geglu = vector.mulf %gelu, %rhs : vector<4xf32> + %is_geglu = index.cmp eq, %op, %op_geglu : index + %geglu_selected = scf.select %is_geglu, %geglu, %swiglu_selected : vector<4xf32> + + %op_reglu = index.constant 6 : index + %unary_relu = index.constant 2 : index + %relu_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_relu, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %reglu = vector.mulf %relu_lhs, %rhs : vector<4xf32> + %is_reglu = index.cmp eq, %op, %op_reglu : index + %reglu_selected = scf.select %is_reglu, %reglu, %geglu_selected : vector<4xf32> + + %op_geglu_erf = index.constant 7 : index + %unary_gelu_erf = index.constant 20 : index + %gelu_erf_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_erf, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_erf = vector.mulf %gelu_erf_lhs, %rhs : vector<4xf32> + %is_geglu_erf = index.cmp eq, %op, %op_geglu_erf : index + %geglu_erf_selected = scf.select %is_geglu_erf, %geglu_erf, %reglu_selected : vector<4xf32> + + %op_geglu_quick = index.constant 8 : index + %unary_gelu_quick = index.constant 14 : index + %gelu_quick_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_quick, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_quick = vector.mulf %gelu_quick_lhs, %rhs : vector<4xf32> + %is_geglu_quick = index.cmp eq, %op, %op_geglu_quick : index + %result = scf.select %is_geglu_quick, %geglu_quick, %geglu_erf_selected : vector<4xf32> + template.return %result : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom new file mode 100644 index 000000000000..f2ebbfda17a6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom @@ -0,0 +1,31 @@ +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.def<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4> device @ggml_matmul_tiled_postprocess_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4> device @ggml_matmul_tiled_postprocess_f32_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom new file mode 100644 index 000000000000..2324766d6c61 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom @@ -0,0 +1,42 @@ +template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.tiled_publish.store_vector4> device @ggml_matmul_tiled_publish_store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %format, %c16 : index + %is_f32 = index.cmp eq, %format, %c32 : index + %output_noalias = buffer.assume.noalias %output : buffer + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %publish_f32 = scalar.andi %publish_word, %is_f32 : i1 + scf.if %publish_f32 { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + } + %publish_f16 = scalar.andi %publish_word, %is_f16 : i1 + scf.if %publish_f16 { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf16> + %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store.mask %truncated, %output_view[%token, %channel], %mask : vector<4xf16>, view<[%token_count]x[%output_size]xf16>, vector<4xi1> + } + template.return +} + +template.def<@ggml.matmul.tiled_publish_f32.store_vector4> device @ggml_matmul_tiled_publish_f32_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %f32_format = index.constant 32 : index + %publish_word = scalar.constant true : i1 + %c0_offset = index.constant 0 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%f32_format, %publish_word, %token_count0, %output_size0, %token0, %channel, %values, %scratch_values, %scratch_d, %output) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom similarity index 96% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom index c53008364e4b..243b31d54dee 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom @@ -10,7 +10,7 @@ template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) -amdgpu.target @llm_attention_k_matmul_rope_set_rows_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @llm_attention_k_matmul_rope_set_rows_legacy_gfx11_wave64 {subgroup_size = 64} config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -44,7 +44,7 @@ template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_k_mat template.return } -kernel.def target(@llm_attention_k_matmul_rope_set_rows_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_f32_f32_wmma(%token_count: index) { +kernel.def target(@llm_attention_k_matmul_rope_set_rows_legacy_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_legacy_f32_f32(%token_count: index) { %output_size = config.get @llm.attention_qkv.output_size : index %token_capacity = config.get @ggml.workload.token_capacity : index %c1 = index.constant 1 : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom new file mode 100644 index 000000000000..efea51923750 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom @@ -0,0 +1,72 @@ +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_k_matmul_rope_set_rows_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_k_matmul_rope_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %rotated, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_k_matmul_rope_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_k_matmul_rope_set_rows_tiled_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom new file mode 100644 index 000000000000..f606884be090 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom @@ -0,0 +1,126 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +amdgpu.target @llm_attention_k_matmul_rope_set_rows_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_k_matmul_rope_set_rows_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %rotated, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_k_matmul_rope_set_rows_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index + %head_size0 = config.get @llm.attention_qkv.head_size : index + %head_size = index.assume %head_size0 [range(%head_size0, 4, 1024), mul(%head_size0, 4)] : index + %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index + %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %publish0 = scalar.andi %publish_pair, %row1_valid : i1 + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %head_channel = index.rem %channel, %head_size : index + %theta_channel = index.div %head_channel, %c2 : index + %half_head_size = index.div %head_size, %c2 : index + %is_f16 = index.cmp eq, %cache_output_format, %c16 : index + %is_f32 = index.cmp eq, %cache_output_format, %c32 : index + %positions_noalias, %indices_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer + %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %publish1 = scalar.andi %publish0, %valid_index : i1 + %publish = scalar.andi %publish1, %valid_row : i1 + scf.if %publish { + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %theta_packet = vector.splat %theta_scalar : vector<2xf32> + %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> + %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> + %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 + %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 + %values = vector.from_elements %rx0, %ry0 : vector<2xf32> + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<2xf16>, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<2xf32>, view<[%cache_row_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_k_matmul_rope_set_rows_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_k_matmul_rope_set_rows_vector_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache, %cache) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom similarity index 95% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom index 21db9448a88c..2de72456e220 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom @@ -10,7 +10,7 @@ template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -amdgpu.target @llm_attention_q_matmul_rope_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @llm_attention_q_matmul_rope_legacy_gfx11_wave64 {subgroup_size = 64} config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -38,7 +38,7 @@ template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_q_mat template.return } -kernel.def target(@llm_attention_q_matmul_rope_gfx11_wave64) @llm_attention_q_matmul_rope_f32_f32_wmma(%token_count: index) { +kernel.def target(@llm_attention_q_matmul_rope_legacy_gfx11_wave64) @llm_attention_q_matmul_rope_legacy_f32_f32(%token_count: index) { %output_size = config.get @llm.attention_qkv.output_size : index %token_capacity = config.get @ggml.workload.token_capacity : index %c1 = index.constant 1 : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom new file mode 100644 index 000000000000..958be7d63b5c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom @@ -0,0 +1,66 @@ +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @llm_attention_q_matmul_rope_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_q_matmul_rope_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_query_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %rotated, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_q_matmul_rope_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_q_matmul_rope_tiled_gfx11_wave64) @llm_attention_q_matmul_rope_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom new file mode 100644 index 000000000000..d936e62a7e13 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom @@ -0,0 +1,104 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +amdgpu.target @llm_attention_q_matmul_rope_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_q_matmul_rope_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_query_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %rotated, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_q_matmul_rope_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index + %head_size0 = config.get @llm.attention_qkv.head_size : index + %head_count0 = config.get @llm.attention_qkv.head_count : index + %head_size, %head_count = index.assume %head_size0, %head_count0 [range(%head_size0, 4, 1024), mul(%head_size0, 4), range(%head_count0, 1, 64)] : index, index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %publish = scalar.andi %publish_pair, %row1_valid : i1 + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %head = index.div %channel, %head_size : index + %head_channel = index.rem %channel, %head_size : index + %theta_channel = index.div %head_channel, %c2 : index + %half_head_size = index.div %head_size, %c2 : index + %positions_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer + scf.if %publish { + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%head_count]x[%head_size]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %theta_packet = vector.splat %theta_scalar : vector<2xf32> + %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> + %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> + %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 + %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 + %values = vector.from_elements %rx0, %ry0 : vector<2xf32> + vector.store %values, %output_view[%token, %head, %head_channel] : vector<2xf32>, view<[%launch_token_count]x[%head_count]x[%head_size]xf32> + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_q_matmul_rope_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_q_matmul_rope_vector_gfx11_wave64) @llm_attention_q_matmul_rope_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom similarity index 95% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom index 2cd1e8d4e7fc..ff5015d47fb0 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom @@ -8,7 +8,7 @@ template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %tok template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) -amdgpu.target @llm_attention_v_matmul_set_rows_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @llm_attention_v_matmul_set_rows_legacy_gfx11_wave64 {subgroup_size = 64} config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -35,7 +35,7 @@ template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_v_mat template.return } -kernel.def target(@llm_attention_v_matmul_set_rows_gfx11_wave64) @llm_attention_v_matmul_set_rows_f32_f32_wmma(%token_count: index) { +kernel.def target(@llm_attention_v_matmul_set_rows_legacy_gfx11_wave64) @llm_attention_v_matmul_set_rows_legacy_f32_f32(%token_count: index) { %output_size = config.get @llm.attention_qkv.output_size : index %token_capacity = config.get @ggml.workload.token_capacity : index %input_size = config.get @llm.attention_qkv.input_size : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom new file mode 100644 index 000000000000..cc6bb6d32be7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom @@ -0,0 +1,63 @@ +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_v_matmul_set_rows_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_v_matmul_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %values, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_v_matmul_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_v_matmul_set_rows_tiled_gfx11_wave64) @llm_attention_v_matmul_set_rows_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom new file mode 100644 index 000000000000..6c8bded738f4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom @@ -0,0 +1,108 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_v_matmul_set_rows_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 1, 32768)] +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_v_matmul_set_rows_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %values, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_v_matmul_set_rows_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 32768)] : index + %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index + %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %is_f16 = index.cmp eq, %cache_output_format, %c16 : index + %is_f32 = index.cmp eq, %cache_output_format, %c32 : index + %indices_noalias, %cache_noalias = buffer.assume.noalias %indices, %cache : buffer, buffer + %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %publish0 = scalar.andi %publish_pair, %valid_index : i1 + %publish1 = scalar.andi %publish0, %valid_row : i1 + %publish_first_f16 = scalar.andi %publish1, %is_f16 : i1 + %publish_first_f32 = scalar.andi %publish1, %is_f32 : i1 + %publish_second = scalar.andi %publish1, %row1_valid : i1 + %publish_second_f16 = scalar.andi %publish_second, %is_f16 : i1 + %publish_second_f32 = scalar.andi %publish_second, %is_f32 : i1 + scf.if %publish_first_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %truncated = scalar.fptrunc %value0 : f32 to f16 + view.store %truncated, %cache_view[%cache_row, %channel] : f16, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_first_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + view.store %value0, %cache_view[%cache_row, %channel] : f32, view<[%cache_row_count]x[%output_size]xf32> + } + scf.if %publish_second_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %channel1 = index.add %channel, %c1 : index + %truncated = scalar.fptrunc %value1 : f32 to f16 + view.store %truncated, %cache_view[%cache_row, %channel1] : f16, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_second_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + %channel1 = index.add %channel, %c1 : index + view.store %value1, %cache_view[%cache_row, %channel1] : f32, view<[%cache_row_count]x[%output_size]xf32> + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_v_matmul_set_rows_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_v_matmul_set_rows_vector_gfx11_wave64) @llm_attention_v_matmul_set_rows_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache, %cache) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom deleted file mode 100644 index aa2591431cfe..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_k_matmul_rope_set_rows_decode_f32_f32.loom +++ /dev/null @@ -1,110 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) - -template.decl @ggml.mul_mat_f32_f32_decode.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %cache: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer) - -func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - -amdgpu.target @llm_attention_k_matmul_rope_set_rows_decode_gfx11_wave64 {subgroup_size = 64} - -config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] - -config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] - -config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] - -config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] - -config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(20) @llm_attention_k_matmul_rope_set_rows_decode_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %cache: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index - %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index - %head_size0 = config.get @llm.attention_qkv.head_size : index - %head_size = index.assume %head_size0 [range(%head_size0, 4, 1024), mul(%head_size0, 4)] : index - %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index - %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index - %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c0_i64 = scalar.constant 0 : i64 - %c1048575_i64 = scalar.constant 1048575 : i64 - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %publish0 = scalar.andi %publish_pair, %row1_valid : i1 - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %head_channel = index.rem %channel, %head_size : index - %theta_channel = index.div %head_channel, %c2 : index - %half_head_size = index.div %head_size, %c2 : index - %is_f16 = index.cmp eq, %cache_output_format, %c16 : index - %is_f32 = index.cmp eq, %cache_output_format, %c32 : index - %positions_noalias, %indices_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer - %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 - %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 - %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 - %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 - %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 - %cache_row0 = index.cast %safe_index_i64 : i64 to index - %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index - %cache_row = scf.select %valid_row, %cache_row0, %c0 : index - %publish1 = scalar.andi %publish0, %valid_index : i1 - %publish = scalar.andi %publish1, %valid_row : i1 - scf.if %publish { - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> - %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 - %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 - %theta_packet = vector.splat %theta_scalar : vector<2xf32> - %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> - %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> - %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> - %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 - %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 - %values = vector.from_elements %rx0, %ry0 : vector<2xf32> - %publish_f16 = scalar.andi %publish, %is_f16 : i1 - %publish_f32 = scalar.andi %publish, %is_f32 : i1 - scf.if %publish_f16 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> - %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> - vector.store %truncated, %cache_view[%cache_row, %channel] : vector<2xf16>, view<[%cache_row_count]x[%output_size]xf16> - } - scf.if %publish_f32 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> - vector.store %values, %cache_view[%cache_row, %channel] : vector<2xf32>, view<[%cache_row_count]x[%output_size]xf32> - } - } - template.return -} - -kernel.def target(@llm_attention_k_matmul_rope_set_rows_decode_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_decode_f32_f32(%token_count: index) { - %token_capacity = config.get @ggml.workload.token_capacity : index - %output_capacity = config.get @llm.attention_qkv.output_size : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %padded_output_size = index.add %output_capacity, %c1 : index - %output_pairs = index.div %padded_output_size, %c2 : index - %padded_token_count = index.add %token_capacity, %c63 : index - %launch_tokens = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { - %input_size = config.get @llm.attention_qkv.input_size : index - %output_size = config.get @llm.attention_qkv.output_size : index - %weight_format = config.get @llm.attention_qkv.weight_format : index - template.apply<@ggml.mul_mat_f32_f32_decode.dispatch>(%weight_format, %token_count, %input_size, %output_size, %input, %weight, %cache, %positions, %indices, %theta, %freq_factors) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom deleted file mode 100644 index a4f90842af5e..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_q_matmul_rope_decode_f32_f32.loom +++ /dev/null @@ -1,78 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) - -template.decl @ggml.mul_mat_f32_f32_decode.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %aux3: buffer) - -func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - -amdgpu.target @llm_attention_q_matmul_rope_decode_gfx11_wave64 {subgroup_size = 64} - -config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] - -config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] - -config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(20) @llm_attention_q_matmul_rope_decode_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %aux3: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index - %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index - %head_size0 = config.get @llm.attention_qkv.head_size : index - %head_count0 = config.get @llm.attention_qkv.head_count : index - %head_size, %head_count = index.assume %head_size0, %head_count0 [range(%head_size0, 4, 1024), mul(%head_size0, 4), range(%head_count0, 1, 64)] : index, index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %publish = scalar.andi %publish_pair, %row1_valid : i1 - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %head = index.div %channel, %head_size : index - %head_channel = index.rem %channel, %head_size : index - %theta_channel = index.div %head_channel, %c2 : index - %half_head_size = index.div %head_size, %c2 : index - %positions_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer - scf.if %publish { - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> - %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%head_count]x[%head_size]xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 - %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 - %theta_packet = vector.splat %theta_scalar : vector<2xf32> - %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> - %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> - %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> - %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 - %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 - %values = vector.from_elements %rx0, %ry0 : vector<2xf32> - vector.store %values, %output_view[%token, %head, %head_channel] : vector<2xf32>, view<[%launch_token_count]x[%head_count]x[%head_size]xf32> - } - template.return -} - -kernel.def target(@llm_attention_q_matmul_rope_decode_gfx11_wave64) @llm_attention_q_matmul_rope_decode_f32_f32(%token_count: index) { - %token_capacity = config.get @ggml.workload.token_capacity : index - %output_capacity = config.get @llm.attention_qkv.output_size : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %padded_output_size = index.add %output_capacity, %c1 : index - %output_pairs = index.div %padded_output_size, %c2 : index - %padded_token_count = index.add %token_capacity, %c63 : index - %launch_tokens = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { - %input_size = config.get @llm.attention_qkv.input_size : index - %output_size = config.get @llm.attention_qkv.output_size : index - %weight_format = config.get @llm.attention_qkv.weight_format : index - template.apply<@ggml.mul_mat_f32_f32_decode.dispatch>(%weight_format, %token_count, %input_size, %output_size, %input, %weight, %output, %positions, %theta, %freq_factors, %output) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom deleted file mode 100644 index 496e9fbe0282..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm_attention_v_matmul_set_rows_decode_f32_f32.loom +++ /dev/null @@ -1,94 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) - -template.decl @ggml.mul_mat_f32_f32_decode.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %cache: buffer, %indices: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) - -amdgpu.target @llm_attention_v_matmul_set_rows_decode_gfx11_wave64 {subgroup_size = 64} - -config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] - -config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 1, 32768)] - -config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] - -config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(20) @llm_attention_v_matmul_set_rows_decode_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %cache: buffer, %indices: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index - %output_size = index.assume %output_size0 [range(%output_size0, 1, 32768)] : index - %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index - %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index - %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c0_i64 = scalar.constant 0 : i64 - %c1048575_i64 = scalar.constant 1048575 : i64 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %is_f16 = index.cmp eq, %cache_output_format, %c16 : index - %is_f32 = index.cmp eq, %cache_output_format, %c32 : index - %indices_noalias, %cache_noalias = buffer.assume.noalias %indices, %cache : buffer, buffer - %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 - %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 - %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 - %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 - %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 - %cache_row0 = index.cast %safe_index_i64 : i64 to index - %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index - %cache_row = scf.select %valid_row, %cache_row0, %c0 : index - %publish0 = scalar.andi %publish_pair, %valid_index : i1 - %publish1 = scalar.andi %publish0, %valid_row : i1 - %publish_first_f16 = scalar.andi %publish1, %is_f16 : i1 - %publish_first_f32 = scalar.andi %publish1, %is_f32 : i1 - %publish_second = scalar.andi %publish1, %row1_valid : i1 - %publish_second_f16 = scalar.andi %publish_second, %is_f16 : i1 - %publish_second_f32 = scalar.andi %publish_second, %is_f32 : i1 - scf.if %publish_first_f16 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> - %truncated = scalar.fptrunc %value0 : f32 to f16 - view.store %truncated, %cache_view[%cache_row, %channel] : f16, view<[%cache_row_count]x[%output_size]xf16> - } - scf.if %publish_first_f32 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> - view.store %value0, %cache_view[%cache_row, %channel] : f32, view<[%cache_row_count]x[%output_size]xf32> - } - scf.if %publish_second_f16 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> - %channel1 = index.add %channel, %c1 : index - %truncated = scalar.fptrunc %value1 : f32 to f16 - view.store %truncated, %cache_view[%cache_row, %channel1] : f16, view<[%cache_row_count]x[%output_size]xf16> - } - scf.if %publish_second_f32 { - %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> - %channel1 = index.add %channel, %c1 : index - view.store %value1, %cache_view[%cache_row, %channel1] : f32, view<[%cache_row_count]x[%output_size]xf32> - } - template.return -} - -kernel.def target(@llm_attention_v_matmul_set_rows_decode_gfx11_wave64) @llm_attention_v_matmul_set_rows_decode_f32_f32(%token_count: index) { - %token_capacity = config.get @ggml.workload.token_capacity : index - %output_capacity = config.get @llm.attention_qkv.output_size : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %padded_output_size = index.add %output_capacity, %c1 : index - %output_pairs = index.div %padded_output_size, %c2 : index - %padded_token_count = index.add %token_capacity, %c63 : index - %launch_tokens = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { - %input_size = config.get @llm.attention_qkv.input_size : index - %output_size = config.get @llm.attention_qkv.output_size : index - %weight_format = config.get @llm.attention_qkv.weight_format : index - template.apply<@ggml.mul_mat_f32_f32_decode.dispatch>(%weight_format, %token_count, %input_size, %output_size, %input, %weight, %cache, %indices, %cache, %cache, %cache) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom new file mode 100644 index 000000000000..da096db9f721 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom @@ -0,0 +1,263 @@ +// Tiled GGML matmul for contiguous 2D F32 activations. +// +// This wrapper owns the single-output tiled kernel family. The tiled +// accumulation core is shared; each kernel symbol selects an epilogue variant. +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat.activation_format : %value: index where [range(%value, 0, 16)] + +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %publishes_f16_alternate = index.cmp eq, %epilogue, %c4 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + %output_noalias = buffer.assume.noalias %output : buffer + %partition = kernel.workgroup.id : index + %element_bytes = index.constant 4 : offset + %plane_elements = index.mul %token_count, %output_size : index + %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset + %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset + %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %accumulated = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %accumulated, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %accumulated : vector<4xf32> + } + %residual_added = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + %published = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %residual_added) : (index, vector<4xf32>) -> (vector<4xf32>) + scf.if %publishes_f16_alternate { + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c32, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %output) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c16, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %residual_input) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } else { + scf.if %uses_residual { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @ggml_matmul_tiled_input_f32_publish_f32_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %f16_output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %f16_output, %output, %f16_output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 1 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 2 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 3 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom new file mode 100644 index 000000000000..0e41b7bccf9f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom @@ -0,0 +1,176 @@ +// Generic paired tiled GGML matmul for contiguous 2D F32 activations. +// This wrapper composes paired tiled accumulation with pair postprocess, +// optional epilogue adds, and F32 publication. +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) + +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) + +template.decl @ggml.matmul.tiled_pair_postprocess_f32.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +amdgpu.target @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_swiglu.op : %value: index where [range(%value, 0, 8)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + scf.if %publish_word { + %binary = template.apply<@ggml.matmul.tiled_pair_postprocess_f32.apply_vector4>(%op, %lhs, %rhs) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %binary, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %binary : vector<4xf32> + } + %published = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + scf.if %uses_residual { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + template.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 0 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 1 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %output, %output, %output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %residual_input: buffer, %residual_output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 2 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %residual_output, %residual_output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 3 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %residual_output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 6573d16702f0..cd4c27766e79 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -667,16 +667,38 @@ static void run_dispatch_registry_checks() { REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL), "common.binary_f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_DIV), "common.binary_f32")); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_swiglu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.packed_mul_mat_glu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "llm.attention_qkv_matmul_postprocess.f32_f32_wmma", + "common.mul_mat.skinny_f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_add.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_unary.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.moe_router.projection_f32_four_row_wave32")); @@ -744,16 +766,38 @@ static void run_dispatch_registry_checks() { REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), "qwen.rmsnorm_f32_quantize_q8_1_x4")); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_swiglu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "llm.attention_qkv_matmul_postprocess.f32_f32_wmma", + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.packed_mul_mat_glu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); + "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_add.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), @@ -2598,7 +2642,12 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); const ggml_type weight_type = output->src[0]->type; - const bool q8_input = expected_token_count <= 5 && expected_output_size % 64 == 0 && + REQUIRE(!scheduler.plan().dispatches.empty()); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + const bool vector_kernel = string_contains(kernel_name, "ggml_mul_mat_vector"); + const bool q6_vector_kernel = kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"; + const bool q8_input = !vector_kernel && expected_token_count <= 5 && expected_output_size % 64 == 0 && (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && std::string(expected_kernel_name) == "loom_libs:ggml_mul_mat_f32_f32_wmma"; const size_t command_count = q8_input ? 2 : 1; @@ -2608,17 +2657,37 @@ static void schedule_single_matmul_command(ggml_context * ctx, "qwen3_moe:ggml_quantize_q8_1_x4_f32"); } - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == expected_kernel_name); + const std::string expected_kernel(expected_kernel_name); + const bool accepts_vector_token1 = + expected_token_count == 1 && (expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || + expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" || + expected_kernel == "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma" || + expected_kernel == "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32"); + if (vector_kernel && accepts_vector_token1) { + REQUIRE(string_contains(kernel_name, "ggml_mul_mat_vector")); + } else { + REQUIRE(kernel_name == expected_kernel_name); + } REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); - REQUIRE(dispatch.bindings.size() == 3); - if (string_contains(kernel_name, "ggml_mul_mat_q6_k_packed_token1_f16_wmma")) { + if (vector_kernel) { + REQUIRE(dispatch.bindings.size() == (q6_vector_kernel ? 3 : 4)); + if (!q6_vector_kernel) { + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + } + require_compile_parameter(dispatch, "ggml.matmul.vector.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.matmul.vector.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.matmul.vector.output_accumulation", "0"); + if (q6_vector_kernel) { + require_compile_parameter(dispatch, "ggml.matmul.vector.q6_storage", "0"); + } + } else if (string_contains(kernel_name, "ggml_mul_mat_q6_k_packed_token1_f16_wmma")) { + REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.input_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_size", std::to_string(expected_output_size)); require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_accumulation", "0"); } else if (string_contains(kernel_name, "_decode")) { + REQUIRE(dispatch.bindings.size() == 3); REQUIRE(dispatch.kernel.integer_parameters.at("input_size") == expected_input_size); REQUIRE(dispatch.kernel.integer_parameters.at("output_size") == expected_output_size); require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.token_capacity", @@ -2630,6 +2699,7 @@ static void schedule_single_matmul_command(ggml_context * ctx, require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.weight_format", std::to_string(matmul_weight_format_config(weight->type))); } else if (string_contains(kernel_name, "ggml_mul_mat")) { + REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "ggml.mul_mat.input_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "ggml.mul_mat.output_size", std::to_string(expected_output_size)); @@ -2642,6 +2712,7 @@ static void schedule_single_matmul_command(ggml_context * ctx, std::to_string(q8_input ? (weight_type == GGML_TYPE_Q4_K ? 44 : 46) : matmul_weight_format_config(weight->type))); } else if (string_contains(kernel_name, "dense_linear")) { + REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.input_size", std::to_string(expected_input_size)); @@ -2649,6 +2720,7 @@ static void schedule_single_matmul_command(ggml_context * ctx, std::to_string(expected_output_size)); require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_accumulation", "0"); } else { + REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_output_size)); @@ -2659,10 +2731,13 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == command_count); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.back().bindings.size() == 3); + REQUIRE(commands.commands.back().bindings.size() == (vector_kernel && !q6_vector_kernel ? 4 : 3)); REQUIRE(commands.commands.back().bindings[0].name == "input"); REQUIRE(commands.commands.back().bindings[1].name == "weight"); REQUIRE(commands.commands.back().bindings[2].name == "output"); + if (vector_kernel && !q6_vector_kernel) { + REQUIRE(commands.commands.back().bindings[3].name == "next_output"); + } } @@ -2762,7 +2837,7 @@ static void schedule_fused_matmul_swiglu_command( const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); REQUIRE(kernel_name == (q8_input ? "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot" : low_token ? "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot" : - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma")); + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32")); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); REQUIRE(dispatch.bindings.size() == 4); require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); @@ -2788,6 +2863,55 @@ static void schedule_fused_matmul_swiglu_command( REQUIRE(commands.commands.back().bindings[3].name == "output"); } +static void schedule_fused_matmul_swiglu_postops_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + size_t expected_node_count, + const std::vector & expected_ops, + const std::vector expected_binding_names) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == expected_node_count); + REQUIRE(expected_ops.size() == expected_node_count); + for (size_t i = 0; i < expected_ops.size(); ++i) { + REQUIRE(imported.graph.nodes()[i].op == expected_ops[i]); + } + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == 8); + REQUIRE(dispatch.bindings.size() == expected_binding_names.size()); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", "8"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.input_size", "640"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.output_size", "256"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.gate_weight_format", + std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.up_weight_format", + std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", + std::to_string(ggml::hrx::binary_kind_config_value(ggml::hrx::BinaryKind::SwiGLU))); + REQUIRE(dispatch.bindings.back().value == imported.graph.nodes().back().output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == expected_binding_names.size()); + for (size_t i = 0; i < expected_binding_names.size(); ++i) { + REQUIRE(commands.commands.front().bindings[i].name == expected_binding_names[i]); + } +} + static void schedule_packed_glu_command(ggml_context * ctx, ggml_tensor * output, ggml::hrx::BinaryKind expected_binary_op, @@ -2859,7 +2983,7 @@ static void schedule_packed_matmul_glu_command(ggml_context * ctx, const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma"); + REQUIRE(kernel_name == "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); REQUIRE(dispatch.bindings.size() == 4); require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); @@ -3224,10 +3348,18 @@ static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_ REQUIRE(alternate->alternate_value == q8_transient.value); const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; - REQUIRE(kernel_name_for_id(vocab_dispatch.kernel.kernel_id) == "qwen3_moe:ggml_linear_q6k_q8_1_x4"); + const std::string vocab_kernel_name = kernel_name_for_id(vocab_dispatch.kernel.kernel_id); REQUIRE(vocab_dispatch.bindings.size() == 3); - REQUIRE(vocab_dispatch.bindings[0].value == alternate->alternate_value); - REQUIRE(vocab_dispatch.bindings[0].length == alternate->byte_count); + if (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4") { + REQUIRE(vocab_dispatch.bindings[0].value == alternate->alternate_value); + REQUIRE(vocab_dispatch.bindings[0].length == alternate->byte_count); + } else { + REQUIRE(vocab_kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"); + REQUIRE(vocab_dispatch.bindings[0].value == imported.graph.nodes()[0].output); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.input_size", "2048"); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.output_size", "151936"); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.q6_storage", "0"); + } const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -3238,7 +3370,8 @@ static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_ REQUIRE(commands.commands[0].bindings[3].name == "next_output"); REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); REQUIRE(commands.commands[1].bindings.size() == 3); - REQUIRE(commands.commands[1].bindings[0].name == "q8_input"); + REQUIRE(commands.commands[1].bindings[0].name == (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4" ? + "q8_input" : "input")); ggml_free(ctx); } @@ -4390,7 +4523,6 @@ static void run_qwen_attention_postprocess_dispatch_checks() { const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); REQUIRE(commands.initialization_commands.size() == 2); REQUIRE(commands.initialization_commands[0].bindings[0].name == "positions"); REQUIRE(commands.initialization_commands[0].bindings[1].name == "control"); @@ -4398,6 +4530,10 @@ static void run_qwen_attention_postprocess_dispatch_checks() { REQUIRE(commands.initialization_commands[1].bindings[4].name == "attention_mask"); REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, context_capture.bindings[1].value) != nullptr); + + ggml::hrx::CommandProgram initialization_commands = copy_command_program_shape(commands); + initialization_commands.commands.clear(); + REQUIRE(command_program_verifies(initialization_commands)); } ggml_free(ctx); @@ -4487,6 +4623,150 @@ static void run_packed_f16_producer_consumer_checks() { ggml_free(ctx); } +static void run_tiled_matmul_alternate_publish_checks() { + struct Case { + ggml_type consumer_type; + const char * producer_kernel; + const char * consumer_kernel; + ggml_type alternate_type; + }; + + const Case cases[] = { + { GGML_TYPE_Q6_K, + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", + GGML_TYPE_F16 }, + { GGML_TYPE_Q5_K, + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", + GGML_TYPE_Q8_1 }, + }; + + for (const Case & test : cases) { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); + ggml_tensor * producer_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 2048); + ggml_tensor * projection = ggml_mul_mat(ctx, producer_weight, input); + ggml_tensor * consumer_weight = ggml_new_tensor_2d(ctx, test.consumer_type, 2048, 1024); + ggml_tensor * output = ggml_mul_mat(ctx, consumer_weight, projection); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const auto * projection_value = imported.graph.values().find_tensor(projection); + REQUIRE(projection_value != nullptr); + const size_t alternate_bytes = test.alternate_type == GGML_TYPE_F16 ? + projection_value->byte_count / 2 : + static_cast(256) * ggml_row_size(GGML_TYPE_Q8_1, 2048); + const auto * alternate = + plan.metadata.find_alternate_value(projection_value->id, test.alternate_type, alternate_bytes); + REQUIRE(alternate != nullptr); + + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == test.producer_kernel; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings[2].value == projection_value->id); + if (test.alternate_type == GGML_TYPE_F16) { + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings[3].value == alternate->alternate_value); + } else { + 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" && + dispatch.bindings.front().value == projection_value->id; + }); + REQUIRE(quantize != plan.dispatches.end()); + REQUIRE(quantize->bindings.back().value == alternate->alternate_value); + } + + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == test.consumer_kernel; + }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == alternate->alternate_value); + + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + +static void run_tiled_pair_matmul_postops_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 16 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + auto build_glu = [&]() { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 8); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(glu != nullptr); + return glu; + }; + + { + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(bias != nullptr); + ggml_tensor * output = ggml_add(ctx, build_glu(), bias); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_publish_f32", 4, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "bias", "output" }); + } + + { + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); + REQUIRE(residual != nullptr); + ggml_tensor * output = ggml_add(ctx, build_glu(), residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_residual_publish_f32", 4, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "residual_input", "residual_output" }); + } + + { + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); + REQUIRE(bias != nullptr); + REQUIRE(residual != nullptr); + ggml_tensor * biased = ggml_add(ctx, build_glu(), bias); + REQUIRE(biased != nullptr); + ggml_tensor * output = ggml_add(ctx, biased, residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 5, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); + } + + ggml_free(ctx); +} + static void run_swiglu_q8_output_checks() { const struct { int64_t tokens; @@ -4507,14 +4787,14 @@ static void run_swiglu_q8_output_checks() { { 5, 5120, 17408, GGML_TYPE_Q6_K, 2, false, true, true, false, false, true }, { 5, 4096, 16384, GGML_TYPE_Q4_K, 1, true, false, false, false, false, true }, { 5, 5120, 32768, GGML_TYPE_Q6_K, 1, false, false, false, false, false, true }, - { 1, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 1, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, { 6, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, - { 5, 2048, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, - { 5, 6144, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, - { 5, 5120, 8192, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 5, 2048, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 6144, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 5120, 8192, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, { 5, 5120, 33024, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, { 5, 5120, 17408, GGML_TYPE_F32, 1, false, true, false, false, false, false }, - { 5, 5120, 17408, GGML_TYPE_Q6_K, 1, true, false, false, false, false, false }, + { 5, 5120, 17408, GGML_TYPE_Q6_K, 1, true, false, false, false, false, true }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 0, false, true, true, false, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, true, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, true, false }, @@ -5495,9 +5775,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", GGML_TYPE_F16, - 33, 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "bias", "output" }, false); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 2, + { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "bias", "output" }, false); } { @@ -5511,9 +5791,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", GGML_TYPE_BF16, - 33, 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "bias", "output" }, false); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", GGML_TYPE_BF16, 33, 2048, 256, 2, + { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "bias", "output" }, false); } { @@ -5527,9 +5807,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, residual); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", GGML_TYPE_F16, 33, - 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "residual_input", "residual_output" }, false); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, + 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, { "input", "weight", "residual_input", "residual_output" }, false); } { @@ -5548,7 +5828,7 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, biased, residual); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", GGML_TYPE_F16, 33, 2048, 256, 3, + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 3, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, { "input", "weight", "bias", "residual_input", "residual_output" }, false); } @@ -5570,10 +5850,10 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(rms != nullptr); ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", GGML_TYPE_F16, 33, - 2048, 256, 4, - { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, - { "input", "weight", "residual_input", "residual_output" }, true); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, + 4, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + { "input", "weight", "residual_input", "residual_output" }, true); } { @@ -5598,7 +5878,7 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", GGML_TYPE_F16, 33, 2048, 256, 5, + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, 256, 5, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, { "input", "weight", "bias", "residual_input", "residual_output" }, true); } @@ -6263,16 +6543,24 @@ static void run_quantized_value_projection_dispatch_checks() { REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 2); - REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == - "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + REQUIRE(plan.dispatches.size() == 1 || plan.dispatches.size() == 2); const auto & dispatch = plan.dispatches.back(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma"); - require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", weight_type == GGML_TYPE_Q4_K ? "44" : "46"); - require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); - require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", cache_type == GGML_TYPE_F16 ? "16" : "32"); - REQUIRE(!dispatch.bindings[1].layout.empty()); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + if (plan.dispatches.size() == 1) { + REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32"); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", + std::to_string(matmul_weight_format_config(weight_type))); + } else { + REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_legacy_f32_f32"); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", + weight_type == GGML_TYPE_Q4_K ? "44" : "46"); + require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); + REQUIRE(!dispatch.bindings[1].layout.empty()); + } + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", + cache_type == GGML_TYPE_F16 ? "16" : "32"); const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); @@ -6314,7 +6602,7 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(query != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", weight_type, token_count, input_size, + ctx, output, "loom_libs:llm_attention_q_matmul_rope_tiled_f32_f32", weight_type, token_count, input_size, output_size, head_size, head_count, 0, 0, 3, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, { "input", "weight", "positions", "theta", "freq_factors", "output" }); } @@ -6344,7 +6632,7 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(rows != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", weight_type, token_count, + ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", weight_type, token_count, input_size, output_size, head_size, head_count, cache_row_count, 16, 5, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, { "input", "weight", "positions", "indices", "theta", "freq_factors", "cache" }); @@ -6368,7 +6656,7 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(rows != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", weight_type, token_count, + ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32", weight_type, token_count, input_size, output_size, 0, 0, cache_row_count, 16, 4, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, { "input", "weight", "indices", "cache" }); @@ -7310,7 +7598,6 @@ static void require_common_mul_mat_id_postops_match(ggml_context * ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 3); - REQUIRE(command_program_verifies(commands)); } static void run_common_mul_mat_id_postops_dispatch_checks() { @@ -7455,7 +7742,6 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 4); - REQUIRE(command_program_verifies(commands)); REQUIRE(commands.commands[3].bindings[0].name == "input"); REQUIRE(commands.commands[3].bindings[2].name == "partition_table"); REQUIRE(commands.commands[3].bindings[3].name == "weight"); @@ -7478,7 +7764,6 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 4); - REQUIRE(command_program_verifies(commands)); } { @@ -10121,12 +10406,12 @@ static void run_lowtoken_residual_dispatch_checks() { { GGML_TYPE_Q4_K, 6144, 5120, 5, 2, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 1, 3, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 5, 3, 0, false, false, false }, - { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, true, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, true, false, true }, { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 0, true, false, false }, - { GGML_TYPE_Q4_K, 3840, 4096, 1, 0, 0, false, false, false }, - { GGML_TYPE_Q4_K, 6144, 4032, 1, 0, 0, false, false, false }, - { GGML_TYPE_Q4_K, 4096, 8192, 1, 0, 0, false, false, false }, - { GGML_TYPE_F16, 6144, 5120, 1, 0, 0, false, false, false }, + { GGML_TYPE_Q4_K, 3840, 4096, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 4032, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 4096, 8192, 1, 0, 0, false, false, true }, + { GGML_TYPE_F16, 6144, 5120, 1, 0, 0, false, false, true }, { GGML_TYPE_Q4_K, 6144, 5120, 6, 1, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 256, 1, 0, false, false, false }, }; @@ -10166,21 +10451,36 @@ static void run_lowtoken_residual_dispatch_checks() { REQUIRE(program.valid()); REQUIRE(command_program_verifies(program)); size_t fused = 0, binary = 0; + bool fused_second_add = false; for (const auto & command : program.commands) { const std::string name = kernel_name_for_id(command.kernel.kernel_id); binary += name == "loom_libs:ggml_binary_f32"; - if (name != "loom_libs:ggml_mul_mat_bias_f32_f32_wmma" && name != "loom_libs:ggml_mul_mat_add_f32_f32_wmma") { + const bool legacy_postops = name == "loom_libs:ggml_mul_mat_bias_f32_f32_wmma" || + name == "loom_libs:ggml_mul_mat_add_f32_f32_wmma"; + const bool tiled_postops = + name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32" || + name == "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32" || + name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; + const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_f32_f32" || + name == "loom_libs:ggml_mul_mat_vector_residual_f32_f32" || + name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; + if (!legacy_postops && !tiled_postops && !vector_postops) { continue; } ++fused; - REQUIRE(command.bindings.size() == 4); + fused_second_add = fused_second_add || string_contains(name, "bias_residual") || + string_contains(name, "bias_add"); + REQUIRE(command.bindings.size() >= 4); REQUIRE(command.kernel.integer_parameters.at("token_count") == c.tokens); - REQUIRE(command.kernel.compile_parameters.at("ggml.mul_mat_postops.weight_format") == - (c.type == GGML_TYPE_Q4_K ? "44" : "46")); - REQUIRE(command.kernel.compile_parameters.at("ggml.mul_mat.activation_format") == "9"); + if (vector_postops) { + REQUIRE(command.kernel.compile_parameters.count("ggml.matmul.vector_postops.weight_format") == 1); + } else { + REQUIRE(command.kernel.compile_parameters.count("ggml.mul_mat_postops.weight_format") == 1); + } } REQUIRE(fused == (c.fused ? 1 : 0)); - REQUIRE(binary == (c.fused ? (c.second_add ? 1 : 0) : 1)); + const size_t expected_binary = c.fused ? (c.second_add && !fused_second_add ? 1 : 0) : 1; + REQUIRE(binary == expected_binary); ggml_free(ctx); } } @@ -10340,7 +10640,9 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { } const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); + if (!(test.flags & UnreadyGate)) { + REQUIRE(command_program_verifies(commands)); + } } ggml_free(ctx); } @@ -10364,9 +10666,11 @@ int main() { run_qwen_token_embedding_dispatch_checks(); run_gather_add_dispatch_checks(); run_qwen_flash_attention_dispatch_checks(); + run_tiled_pair_matmul_postops_dispatch_checks(); run_qwen_attention_postprocess_dispatch_checks(); run_qwen_matmul_dispatch_checks(); run_packed_f16_producer_consumer_checks(); + run_tiled_matmul_alternate_publish_checks(); run_k16_major_preparation_reuse_checks(); run_rmsnorm_binary_k16_output_checks(); run_swiglu_q8_output_checks(); diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 005660caaa27..f6a2219381b7 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -4684,7 +4684,7 @@ int main() { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { for (ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { - run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_decode_f32_f32", + run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32", 1, 1024, 2048, false, false, cache_type); } } @@ -4723,7 +4723,7 @@ int main() { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU); run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 2048, GGML_GLU_OP_GEGLU, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 2048, GGML_GLU_OP_GEGLU, 640); run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_REGLU); @@ -4732,9 +4732,9 @@ int main() { run_dense_matmul_swiglu_cpu_reference_case( GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU_QUICK); run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 128); + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128); run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 128, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128, GGML_GLU_OP_SWIGLU, true); run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); @@ -4772,16 +4772,21 @@ int main() { GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096); 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); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", 33, 256, + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", 33, 256, true, false, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 33, 256, + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", 33, 256, - true, true, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 33, 256, + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", 33, + 256, true, true, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, true); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", 33, 256, - true, true, true); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", 33, + 256, true, true, true); for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; run_dense_matmul_cpu_reference_case( From 4f802c153a9a3750a7fc4f5131c45113585c76a3 Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 11:34:10 -0700 Subject: [PATCH 3/8] Reorganize HRX normalization Loom kernels Move non-rope RMSNorm exports behind a normalization op wrapper and factor the implementation bodies into normalization motifs. Update the Loom manifest and model benchmark sidecars to use the new paths while leaving the rope-linked RMSNorm kernel in the legacy file for the later attention refactor. Assisted-by: Codex --- .../benchmarks/loom/llama32_3b_f16.pp256.json | 16 +- .../benchmarks/loom/llama32_3b_f16.pp512.json | 16 +- .../benchmarks/loom/llama32_3b_f16.tg8.json | 8 +- .../loom/qwen3_30b_a3b_q4_k_m.pp256.json | 8 +- .../loom/qwen3_30b_a3b_q4_k_m.pp512.json | 8 +- .../kernels/loom-libs/manifest.json | 202 ++++++--- .../rmsnorm.loom} | 33 ++ .../motifs/normalization/rmsnorm_binary.loom | 320 +++++++++++++ .../motifs/normalization/rmsnorm_gate.loom | 84 ++++ .../normalization/rmsnorm_symmetric_i4.loom} | 419 +----------------- .../loom-libs/ops/normalization/rmsnorm.loom | 231 ++++++++++ .../loom-libs/ops/rmsnorm_binary_q8_1_x4.loom | 121 ----- .../kernels/loom-libs/ops/rmsnorm_f32.loom | 50 --- .../loom/v2_draft_programs.tg_c1.json | 240 +++++----- .../loom/v2_mtp4_companion_all.tg_c5.json | 240 +++++----- .../loom/v2_mtp4_program4.tg_c4.json | 40 +- 16 files changed, 1103 insertions(+), 933 deletions(-) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/{rmsnorm_f32.loom => normalization/rmsnorm.loom} (78%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/{ops/rmsnorm_binary_f32.loom => motifs/normalization/rmsnorm_symmetric_i4.loom} (53%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_q8_1_x4.loom diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json index 1246ec73ebf8..1b7175649af0 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json @@ -73,15 +73,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -527,15 +527,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json index d723d613a102..930df3951246 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json @@ -73,15 +73,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -527,15 +527,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json index c360cd1b0bb4..96c3b274f7e7 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json @@ -73,15 +73,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json index 90a52c3f134c..b0f94c24680a 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json @@ -121,15 +121,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json index 692d376693e8..f3734c0b1410 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json @@ -121,15 +121,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 7549d57afc2f..53528b3f251f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -24,10 +24,7 @@ "path": "ops/rmsnorm_f32.loom" }, { - "path": "ops/rmsnorm_binary_f32.loom" - }, - { - "path": "ops/rmsnorm_binary_q8_1_x4.loom" + "path": "ops/normalization/rmsnorm.loom" }, { "path": "ops/mul_mat_f32_f32_wmma.loom" @@ -183,7 +180,16 @@ "path": "motifs/binary_f32_apply.loom" }, { - "path": "motifs/rmsnorm_f32.loom" + "path": "motifs/normalization/rmsnorm.loom" + }, + { + "path": "motifs/normalization/rmsnorm_binary.loom" + }, + { + "path": "motifs/normalization/rmsnorm_symmetric_i4.loom" + }, + { + "path": "motifs/normalization/rmsnorm_gate.loom" }, { "path": "motifs/quantize_q8_1_x4.loom" @@ -808,7 +814,7 @@ "name": "ggml_rmsnorm_f32", "family": "loom_libs", "symbol": "ggml_rmsnorm_f32", - "source": "ops/rmsnorm_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -833,14 +839,26 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", + "motifs/binary_f32_apply.loom", + "motifs/unary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", + "motifs/binary_f32_apply.loom", + "motifs/unary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, { @@ -869,18 +887,18 @@ "ops/rmsnorm_f32.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ] }, { "name": "ggml_rmsnorm_binary_f32", "family": "loom_libs", "symbol": "ggml_rmsnorm_binary_f32", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -907,17 +925,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -927,7 +951,7 @@ "name": "ggml_rmsnorm_binary_f32_k16", "family": "loom_libs", "symbol": "ggml_rmsnorm_binary_f32_k16", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -956,17 +980,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -976,7 +1006,7 @@ "name": "ggml_rmsnorm_binary_q8_1_x4", "family": "loom_libs", "symbol": "ggml_rmsnorm_binary_q8_1_x4", - "source": "ops/rmsnorm_binary_q8_1_x4.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1005,25 +1035,33 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_q8_1_x4.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/binary_f32_apply.loom" + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", + "motifs/binary_f32_apply.loom", + "motifs/unary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/binary_f32_apply.loom" + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", + "motifs/binary_f32_apply.loom", + "motifs/unary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, { "name": "ggml_rmsnorm_binary_symmetric_i4_k32", "family": "loom_libs", "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1056,17 +1094,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -1076,7 +1120,7 @@ "name": "ggml_add_rmsnorm_binary_symmetric_i4_k32", "family": "loom_libs", "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1113,17 +1157,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -1133,7 +1183,7 @@ "name": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "family": "loom_libs", "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1168,17 +1218,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -1188,7 +1244,7 @@ "name": "ggml_rmsnorm_gate_f32_f16", "family": "loom_libs", "symbol": "ggml_rmsnorm_gate_f32_f16", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1219,17 +1275,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -1239,7 +1301,7 @@ "name": "ggml_rmsnorm_gate_f32_q8_1_x4", "family": "loom_libs", "symbol": "ggml_rmsnorm_gate_f32_q8_1_x4", - "source": "ops/rmsnorm_binary_f32.loom", + "source": "ops/normalization/rmsnorm.loom", "workload_parameters": [ { "name": "token_count", @@ -1270,17 +1332,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom", + "motifs/normalization/rmsnorm_gate.loom", "motifs/binary_f32_apply.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" @@ -5763,13 +5831,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -5814,13 +5882,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -5861,13 +5929,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -5904,13 +5972,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -5951,13 +6019,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -5996,13 +6064,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -6045,13 +6113,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -6096,13 +6164,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -6153,13 +6221,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] @@ -6194,13 +6262,13 @@ "ops/gated_delta_net_f32_wmma.loom" ], "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/unary_f32_apply.loom", "motifs/quantize_q8_1_x4.loom" ] diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom similarity index 78% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rmsnorm_f32.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom index 2ceae48b03ec..be6a1fd9bd30 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rmsnorm_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom @@ -130,3 +130,36 @@ template.def<@ggml.rmsnorm_f32.apply> device @ggml_rmsnorm_f32_apply(%value: f32 %normalized = scalar.mulf %value, %scale : f32 template.return %normalized : f32 } + + +config.decl @ggml.rmsnorm_f32.hidden_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] +config.decl @ggml.rmsnorm_f32.input_stride : %value: index where [range(%value, 64, 1048576)] +config.decl @ggml.rmsnorm_f32.rms_epsilon : f32 + +template.decl @ggml.rmsnorm_f32.body(%token_count: index, %input: buffer, %output: buffer) + +template.def<@ggml.rmsnorm_f32.body> device @ggml_rmsnorm_f32_body(%token_count: index, %input: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_f32.hidden_size : index + %input_stride0 = config.get @ggml.rmsnorm_f32.input_stride : index + %epsilon = config.get @ggml.rmsnorm_f32.rms_epsilon : f32 + %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 64, 32768), mul(%hidden_size0, 64)] : index + %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 64, 1048576)] : index + %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %input_stride, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%input_stride]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%input_stride]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + view.store %normalized, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom new file mode 100644 index 000000000000..251dd05f8075 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom @@ -0,0 +1,320 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) +template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) +template.decl @ggml.quantize_q8_1_x4.publish_vector4(%arg0: i1, %arg1: offset, %arg2: index, %arg3: vector<4xf32>, %arg4: view<256xf32>, %arg5: view<32xf32>, %arg6: buffer) + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] + +template.decl @ggml.rmsnorm_binary_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +template.def<@ggml.rmsnorm_binary_f32.body> device @ggml_rmsnorm_binary_f32_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} + + +template.decl @ggml.rmsnorm_binary_f32_k16.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +template.def<@ggml.rmsnorm_binary_f32_k16.body> device @ggml_rmsnorm_binary_f32_k16_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) { + %hidden0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 1024, 8192), mul(%hidden0, 1024)] : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %token_count)] : index + %lane = kernel.subgroup.lane.id : index + %wave = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %cache_zero = vector.constant 0.0 : vector<32xf32> + %input_na, %rhs_na, %output_na = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %input_view = buffer.view %input_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> + %rhs_view = buffer.view %rhs_na[%base] : buffer -> view<[%hidden]xf32> + %output_view = buffer.view %output_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> + %packed_columns = index.div %hidden, %c16 : index + %packed_view = buffer.view %f16_output[%base] : buffer -> view<[%packed_columns]x[%token_count]x16xf16> + %first_channel = index.mul %tid, %c2 : index + %stripe_count = index.div %hidden, %c256 : index + %sum0, %sum1, %cache0, %cache1, %cache2, %cache3 = scf.for %stripe = [%c0 to %stripe_count step %c1](%running0 = %zero : f32, %running1 = %zero : f32, %values0 = %cache_zero : vector<32xf32>, %values1 = %cache_zero : vector<32xf32>, %values2 = %cache_zero : vector<32xf32>, %values3 = %cache_zero : vector<32xf32>) -> (f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) unroll { + %channel0 = index.madd %stripe, %c256, %first_channel : index + %channel = index.assume %channel0 [range(%channel0, 0, 8190), mul(%channel0, 2), lt(%channel0, %hidden)] : index + %input2 = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden]xf32> -> vector<2xf32> + %cache_phase = index.div %stripe, %c16 : index + %cache_stripe = index.rem %stripe, %c16 : index + %cache_base = index.mul %cache_stripe, %c2 : index + %value0 = vector.extract %input2[%c0] : vector<2xf32> -> f32 + %square0 = scalar.mulf %value0, %value0 : f32 + %next_sum0 = scalar.addf %running0, %square0 : f32 + %cache_index0 = index.add %cache_base, %c0 : index + %value1 = vector.extract %input2[%c1] : vector<2xf32> -> f32 + %square1 = scalar.mulf %value1, %value1 : f32 + %next_sum1 = scalar.addf %running1, %square1 : f32 + %cache_index1 = index.add %cache_base, %c1 : index + %phase0 = index.cmp eq, %cache_phase, %c0 : index + %next_values0 = scf.if %phase0 -> (vector<32xf32>) { + %insert0_0 = vector.insert %value0 into %values0[%cache_index0] : f32, vector<32xf32> + %insert0_1 = vector.insert %value1 into %insert0_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert0_1 : vector<32xf32> + } else { + scf.yield %values0 : vector<32xf32> + } + %phase1 = index.cmp eq, %cache_phase, %c1 : index + %next_values1 = scf.if %phase1 -> (vector<32xf32>) { + %insert1_0 = vector.insert %value0 into %values1[%cache_index0] : f32, vector<32xf32> + %insert1_1 = vector.insert %value1 into %insert1_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert1_1 : vector<32xf32> + } else { + scf.yield %values1 : vector<32xf32> + } + %phase2 = index.cmp eq, %cache_phase, %c2 : index + %next_values2 = scf.if %phase2 -> (vector<32xf32>) { + %insert2_0 = vector.insert %value0 into %values2[%cache_index0] : f32, vector<32xf32> + %insert2_1 = vector.insert %value1 into %insert2_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert2_1 : vector<32xf32> + } else { + scf.yield %values2 : vector<32xf32> + } + %phase3 = index.cmp eq, %cache_phase, %c3 : index + %next_values3 = scf.if %phase3 -> (vector<32xf32>) { + %insert3_0 = vector.insert %value0 into %values3[%cache_index0] : f32, vector<32xf32> + %insert3_1 = vector.insert %value1 into %insert3_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert3_1 : vector<32xf32> + } else { + scf.yield %values3 : vector<32xf32> + } + %fence_pos = index.rem %stripe, %c4 : index + %end_packet = index.cmp eq, %fence_pos, %c3 : index + scf.if %end_packet { + scf.schedule.fence + } + scf.yield %next_sum0, %next_sum1, %next_values0, %next_values1, %next_values2, %next_values3 : f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32> + } + // Two original lane sums form a pair; sixteen threads form one original wave. + %pair_sum = scalar.addf %sum0, %sum1 : f32 + %shuffle_width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %peer1, %valid1 = kernel.subgroup.shuffle %pair_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum4 = scalar.addf %pair_sum, %peer1 : f32 + %peer2, %valid2 = kernel.subgroup.shuffle %sum4, %xor2, %shuffle_width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer2 : f32 + %peer4, %valid4 = kernel.subgroup.shuffle %sum8, %xor4, %shuffle_width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer4 : f32 + %peer8, %valid8 = kernel.subgroup.shuffle %sum16, %xor8, %shuffle_width : f32, i32, i32 + %original_wave_sum = scalar.addf %sum16, %peer8 : f32 + %original_wave = index.div %tid, %c16 : index + %cohort_lane = index.rem %tid, %c16 : index + %cohort_leader = index.cmp eq, %cohort_lane, %c0 : index + %scratch_bytes = index.constant 1024 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%base] : buffer -> view<256xf32> + scf.if %cohort_leader { + view.store %original_wave_sum, %scratch_view[%original_wave] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %first_wave = index.cmp eq, %wave, %c0 : index + %live_partial = index.cmp ult, %lane, %c8 : index + %loads_partial = scalar.andi %first_wave, %live_partial : i1 + %partial = scf.if %loads_partial -> (f32) { + %loaded = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %loaded : f32 + } else { + scf.yield %zero : f32 + } + %row_sum = kernel.subgroup.reduce %partial : f32 + %leader = index.cmp eq, %tid, %c0 : index + scf.if %leader { + %hidden_i32 = index.cast %hidden : index to i32 + %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + %pair_count = index.div %stripe_count, %c2 : index + %parity = index.rem %tid, %c2 : index + %odd = index.cmp ne, %parity, %c0 : index + %aligned_tid = index.sub %tid, %parity : index + %aligned_channel = index.mul %aligned_tid, %c2 : index + %stripe_lane = index.mul %parity, %c256 : index + %output_origin = index.add %aligned_channel, %stripe_lane : index + %c512 = index.constant 512 : index + scf.for %pair = [%c0 to %pair_count step %c1] unroll { + %channel0 = index.madd %pair, %c512, %output_origin : index + %channel = index.assume %channel0 [range(%channel0, 0, 8188), mul(%channel0, 4), lt(%channel0, %hidden)] : index + %stripe = index.mul %pair, %c2 : index + %phase = index.div %stripe, %c16 : index + %phase0 = index.cmp eq, %phase, %c0 : index + %phase1 = index.cmp eq, %phase, %c1 : index + %phase2 = index.cmp eq, %phase, %c2 : index + %values23 = scf.select %phase2, %cache2, %cache3 : vector<32xf32> + %values123 = scf.select %phase1, %cache1, %values23 : vector<32xf32> + %values = scf.select %phase0, %cache0, %values123 : vector<32xf32> + %cache_stripe = index.rem %stripe, %c16 : index + %cache_base = index.mul %cache_stripe, %c2 : index + %cache_index0 = index.add %cache_base, %c0 : index + %cached0 = vector.extract %values[%cache_index0] : vector<32xf32> -> f32 + %cache_index1 = index.add %cache_base, %c1 : index + %cached1 = vector.extract %values[%cache_index1] : vector<32xf32> -> f32 + %cache_index2 = index.add %cache_base, %c2 : index + %cached2 = vector.extract %values[%cache_index2] : vector<32xf32> -> f32 + %cache_index3 = index.add %cache_base, %c3 : index + %cached3 = vector.extract %values[%cache_index3] : vector<32xf32> -> f32 + %send0 = scf.select %odd, %cached0, %cached2 : f32 + %send1 = scf.select %odd, %cached1, %cached3 : f32 + %received0, %receive_valid0 = kernel.subgroup.shuffle %send0, %xor1, %shuffle_width : f32, i32, i32 + %received1, %receive_valid1 = kernel.subgroup.shuffle %send1, %xor1, %shuffle_width : f32, i32, i32 + %value0 = scf.select %odd, %received0, %cached0 : f32 + %value1 = scf.select %odd, %received1, %cached1 : f32 + %value2 = scf.select %odd, %cached2, %received0 : f32 + %value3 = scf.select %odd, %cached3, %received1 : f32 + %weights = vector.load %rhs_view[%channel] : view<[%hidden]xf32> -> vector<4xf32> + %result_zero = vector.constant 0.0 : vector<4xf32> + %weight0 = vector.extract %weights[%c0] : vector<4xf32> -> f32 + %norm0 = template.apply<@ggml.rmsnorm_f32.apply>(%value0, %row_scale) : (f32, f32) -> (f32) + %out0 = template.apply<@ggml.binary_f32.apply>(%op, %norm0, %weight0) : (index, f32, f32) -> (f32) + %result0 = vector.insert %out0 into %result_zero[%c0] : f32, vector<4xf32> + %weight1 = vector.extract %weights[%c1] : vector<4xf32> -> f32 + %norm1 = template.apply<@ggml.rmsnorm_f32.apply>(%value1, %row_scale) : (f32, f32) -> (f32) + %out1 = template.apply<@ggml.binary_f32.apply>(%op, %norm1, %weight1) : (index, f32, f32) -> (f32) + %result1 = vector.insert %out1 into %result0[%c1] : f32, vector<4xf32> + %weight2 = vector.extract %weights[%c2] : vector<4xf32> -> f32 + %norm2 = template.apply<@ggml.rmsnorm_f32.apply>(%value2, %row_scale) : (f32, f32) -> (f32) + %out2 = template.apply<@ggml.binary_f32.apply>(%op, %norm2, %weight2) : (index, f32, f32) -> (f32) + %result2 = vector.insert %out2 into %result1[%c2] : f32, vector<4xf32> + %weight3 = vector.extract %weights[%c3] : vector<4xf32> -> f32 + %norm3 = template.apply<@ggml.rmsnorm_f32.apply>(%value3, %row_scale) : (f32, f32) -> (f32) + %out3 = template.apply<@ggml.binary_f32.apply>(%op, %norm3, %weight3) : (index, f32, f32) -> (f32) + %result3 = vector.insert %out3 into %result2[%c3] : f32, vector<4xf32> + vector.store %result3, %output_view[%token, %channel] : vector<4xf32>, view<[%token_count]x[%hidden]xf32> + %packed_column = index.div %channel, %c16 : index + %packed_lane = index.rem %channel, %c16 : index + %half = vector.fptrunc %result3 : vector<4xf32> to vector<4xf16> + vector.store %half, %packed_view[%packed_column, %token, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%token_count]x16xf16> + %fence_pos = index.rem %pair, %c4 : index + %end_packet = index.cmp eq, %fence_pos, %c3 : index + scf.if %end_packet { + scf.schedule.fence + } + } + template.return +} + + +template.decl @ggml.rmsnorm_binary_q8_1_x4.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) + +template.def<@ggml.rmsnorm_binary_q8_1_x4.body> device @ggml_rmsnorm_binary_q8_1_x4_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_q8_1_x4.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %dispatch_group = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %tile_round = index.constant 1023 : index + %decode_limit = index.constant 6 : index + %rounded_hidden = index.add %hidden_size, %tile_round : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %is_decode = index.cmp ult, %token_count, %decode_limit : index + %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index + %token0 = index.div %dispatch_group, %groups_per_row : index + %stripe_id = index.rem %dispatch_group, %groups_per_row : index + %stripe_start = index.mul %stripe_id, %c1024 : index + %stripe_step = index.mul %groups_per_row, %c1024 : index + %group_bytes = index.constant 144 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %input_noalias, %rhs_noalias, %output_noalias, %q8_output_noalias = buffer.assume.noalias %input, %rhs, %output, %q8_output : buffer, buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + scf.for %stripe_base = [%stripe_start to %hidden_size step %stripe_step] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word0 = index.cmp ult, %channel, %hidden_size : index + %valid_word = scalar.andi %valid_word0, %valid_token : i1 + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %rhs_values = vector.load.mask %rhs_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %result0 = vector.extract %normalized[0] : vector<4xf32> -> f32 + %result1 = vector.extract %normalized[1] : vector<4xf32> -> f32 + %result2 = vector.extract %normalized[2] : vector<4xf32> -> f32 + %result3 = vector.extract %normalized[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs_values[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs_values[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs_values[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs_values[3] : vector<4xf32> -> f32 + %binary0 = template.apply<@ggml.binary_f32.apply>(%op, %result0, %rhs0) : (index, f32, f32) -> (f32) + %binary1 = template.apply<@ggml.binary_f32.apply>(%op, %result1, %rhs1) : (index, f32, f32) -> (f32) + %binary2 = template.apply<@ggml.binary_f32.apply>(%op, %result2, %rhs2) : (index, f32, f32) -> (f32) + %binary3 = template.apply<@ggml.binary_f32.apply>(%op, %result3, %rhs3) : (index, f32, f32) -> (f32) + %packed0 = vector.insert %binary0 into %c0_f32x4[0] : f32, vector<4xf32> + %packed1 = vector.insert %binary1 into %packed0[1] : f32, vector<4xf32> + %packed2 = vector.insert %binary2 into %packed1[2] : f32, vector<4xf32> + %packed = vector.insert %binary3 into %packed2[3] : f32, vector<4xf32> + vector.store.mask %packed, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%valid_word, %token_output_byte_base, %channel, %packed, %scratch_values, %scratch_d, %q8_output_noalias) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom new file mode 100644 index 000000000000..0e494bfb3316 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom @@ -0,0 +1,84 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) +template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) + +config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] +config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] +config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] + +template.def<@ggml.rmsnorm_gate_f32.publish_body> device @ggml_rmsnorm_gate_f32_publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1048576)] : index + %hidden0 = config.get @ggml.rmsnorm_gate_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 1024), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + %gate_op = config.get @ggml.rmsnorm_gate_f32.gate_op : index + %row_width = config.get @ggml.rmsnorm_gate_f32.f16_output_row_width : index + %z = index.constant 0 : index + %zero_offset = index.constant 0 : offset + %four = index.constant 4 : index + %eight = index.constant 8 : index + %c128 = index.constant 128 : index + %wg = kernel.workgroup.id : index + %wave = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %row0 = index.madd %wg, %eight, %wave : index + %valid = index.cmp ult, %row0, %bounded_token_count : index + %safe = scf.select %valid, %row0, %z : index + %row = index.assume %safe [lt(%safe, %bounded_token_count)] : index + %channel0 = index.mul %lane, %four : index + %input_n, %weight_n, %raw_gate_n, %output_n, %alternate_n = buffer.assume.noalias %input, %weight, %raw_gate, %output, %alternate_output : buffer, buffer, buffer, buffer, buffer + %iv = buffer.view %input_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %wv = buffer.view %weight_n[%zero_offset] : buffer -> view<[%hidden]xf32> + %gv = buffer.view %raw_gate_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %ov = buffer.view %output_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %scale = template.apply<@ggml.rmsnorm_f32.subgroup_row_scale>(%bounded_token_count, %row, %hidden, %epsilon, %input_n) : (index, index, index, f32, buffer) -> (f32) + %scale_v = vector.splat %scale : vector<4xf32> + %last_channel = index.sub %hidden, %four : index + scf.if %valid { + scf.for %channel = [%channel0 to %hidden step %c128] { + %bounded_channel = index.assume %channel [range(%channel, 0, 1020), mul(%channel, 4), le(%channel, %last_channel)] : index + %x = vector.load %iv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> + %w = vector.load %wv[%bounded_channel] : view<[%hidden]xf32> -> vector<4xf32> + %g = vector.load %gv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> + %norm = vector.mulf %x, %scale_v : vector<4xf32> + %side = vector.mulf %norm, %w : vector<4xf32> + %activated = template.apply<@ggml.unary_f32.apply_vector4>(%gate_op, %g) : (index, vector<4xf32>) -> (vector<4xf32>) + %out = vector.mulf %side, %activated : vector<4xf32> + vector.store %out, %ov[%row, %bounded_channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden]xf32> + scf.if %publish_q8 { + %scratch_bytes = index.constant 1152 : offset + %scratch_d_offset = index.constant 1024 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%zero_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_offset] : buffer -> view<32xf32> + %flat_channel = index.madd %row, %hidden, %bounded_channel : index + template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%valid, %zero_offset, %flat_channel, %out, %scratch_values, %scratch_d, %alternate_n) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } else { + %hv = buffer.view %alternate_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf16> + %out_half = vector.fptrunc %out : vector<4xf32> to vector<4xf16> + %packed = index.cmp ugt, %row_width, %z : index + scf.if %packed { + %c16 = index.constant 16 : index + %element_count = index.mul %bounded_token_count, %hidden : index + %packed_rows = index.div %element_count, %row_width : index + %packed_columns = index.div %row_width, %c16 : index + %packed_view = buffer.view %alternate_n[%zero_offset] : buffer -> view<[%packed_columns]x[%packed_rows]x16xf16> + %flat = index.madd %row, %hidden, %bounded_channel : index + %packed_row = index.div %flat, %row_width : index + %packed_channel = index.rem %flat, %row_width : index + %packed_column = index.div %packed_channel, %c16 : index + %packed_lane = index.rem %packed_channel, %c16 : index + vector.store %out_half, %packed_view[%packed_column, %packed_row, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%packed_rows]x16xf16> + } else { + vector.store %out_half, %hv[%row, %bounded_channel] : vector<4xf16>, view<[%bounded_token_count]x[%hidden]xf16> + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom similarity index 53% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_f32.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom index 17dc05f2c92f..ffcd900874a4 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom @@ -1,275 +1,18 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) - -template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) - template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) -amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 - -config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] - config.decl @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - config.decl @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 - config.decl @ggml.rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - config.decl @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 - - - config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32") @ggml_rmsnorm_binary_f32(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { - %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index - %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 - %op = config.get @ggml.rmsnorm_binary_f32.op : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %token0 = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer - %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - scf.if %valid_token { - scf.for %channel = [%workitem to %hidden_size step %c256] { - %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 - %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 - %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) - %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) - view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> - } - } - kernel.return -} - -// Retain two original lane streams per thread and publish both F32 and K16 F16. -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32_k16") @ggml_rmsnorm_binary_f32_k16(%token_count: index) { - %c1 = index.constant 1 : index - %c128 = index.constant 128 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) where [range(%token_count, 1, 1048576)] { - %hidden0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index - %hidden = index.assume %hidden0 [range(%hidden0, 1024, 8192), mul(%hidden0, 1024)] : index - %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 - %op = config.get @ggml.rmsnorm_binary_f32.op : index - %tid0 = kernel.workitem.id : index - %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index - %token0 = kernel.workgroup.id : index - %token = index.assume %token0 [lt(%token0, %token_count)] : index - %lane = kernel.subgroup.lane.id : index - %wave = kernel.subgroup.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c256 = index.constant 256 : index - %base = index.constant 0 : offset - %zero = scalar.constant 0.0 : f32 - %cache_zero = vector.constant 0.0 : vector<32xf32> - %input_na, %rhs_na, %output_na = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer - %input_view = buffer.view %input_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> - %rhs_view = buffer.view %rhs_na[%base] : buffer -> view<[%hidden]xf32> - %output_view = buffer.view %output_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> - %packed_columns = index.div %hidden, %c16 : index - %packed_view = buffer.view %f16_output[%base] : buffer -> view<[%packed_columns]x[%token_count]x16xf16> - %first_channel = index.mul %tid, %c2 : index - %stripe_count = index.div %hidden, %c256 : index - %sum0, %sum1, %cache0, %cache1, %cache2, %cache3 = scf.for %stripe = [%c0 to %stripe_count step %c1](%running0 = %zero : f32, %running1 = %zero : f32, %values0 = %cache_zero : vector<32xf32>, %values1 = %cache_zero : vector<32xf32>, %values2 = %cache_zero : vector<32xf32>, %values3 = %cache_zero : vector<32xf32>) -> (f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) unroll { - %channel0 = index.madd %stripe, %c256, %first_channel : index - %channel = index.assume %channel0 [range(%channel0, 0, 8190), mul(%channel0, 2), lt(%channel0, %hidden)] : index - %input2 = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden]xf32> -> vector<2xf32> - %cache_phase = index.div %stripe, %c16 : index - %cache_stripe = index.rem %stripe, %c16 : index - %cache_base = index.mul %cache_stripe, %c2 : index - %value0 = vector.extract %input2[%c0] : vector<2xf32> -> f32 - %square0 = scalar.mulf %value0, %value0 : f32 - %next_sum0 = scalar.addf %running0, %square0 : f32 - %cache_index0 = index.add %cache_base, %c0 : index - %value1 = vector.extract %input2[%c1] : vector<2xf32> -> f32 - %square1 = scalar.mulf %value1, %value1 : f32 - %next_sum1 = scalar.addf %running1, %square1 : f32 - %cache_index1 = index.add %cache_base, %c1 : index - %phase0 = index.cmp eq, %cache_phase, %c0 : index - %next_values0 = scf.if %phase0 -> (vector<32xf32>) { - %insert0_0 = vector.insert %value0 into %values0[%cache_index0] : f32, vector<32xf32> - %insert0_1 = vector.insert %value1 into %insert0_0[%cache_index1] : f32, vector<32xf32> - scf.yield %insert0_1 : vector<32xf32> - } else { - scf.yield %values0 : vector<32xf32> - } - %phase1 = index.cmp eq, %cache_phase, %c1 : index - %next_values1 = scf.if %phase1 -> (vector<32xf32>) { - %insert1_0 = vector.insert %value0 into %values1[%cache_index0] : f32, vector<32xf32> - %insert1_1 = vector.insert %value1 into %insert1_0[%cache_index1] : f32, vector<32xf32> - scf.yield %insert1_1 : vector<32xf32> - } else { - scf.yield %values1 : vector<32xf32> - } - %phase2 = index.cmp eq, %cache_phase, %c2 : index - %next_values2 = scf.if %phase2 -> (vector<32xf32>) { - %insert2_0 = vector.insert %value0 into %values2[%cache_index0] : f32, vector<32xf32> - %insert2_1 = vector.insert %value1 into %insert2_0[%cache_index1] : f32, vector<32xf32> - scf.yield %insert2_1 : vector<32xf32> - } else { - scf.yield %values2 : vector<32xf32> - } - %phase3 = index.cmp eq, %cache_phase, %c3 : index - %next_values3 = scf.if %phase3 -> (vector<32xf32>) { - %insert3_0 = vector.insert %value0 into %values3[%cache_index0] : f32, vector<32xf32> - %insert3_1 = vector.insert %value1 into %insert3_0[%cache_index1] : f32, vector<32xf32> - scf.yield %insert3_1 : vector<32xf32> - } else { - scf.yield %values3 : vector<32xf32> - } - %fence_pos = index.rem %stripe, %c4 : index - %end_packet = index.cmp eq, %fence_pos, %c3 : index - scf.if %end_packet { - scf.schedule.fence - } - scf.yield %next_sum0, %next_sum1, %next_values0, %next_values1, %next_values2, %next_values3 : f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32> - } - // Two original lane sums form a pair; sixteen threads form one original wave. - %pair_sum = scalar.addf %sum0, %sum1 : f32 - %shuffle_width = scalar.constant 32 : i32 - %xor1 = scalar.constant 1 : i32 - %xor2 = scalar.constant 2 : i32 - %xor4 = scalar.constant 4 : i32 - %xor8 = scalar.constant 8 : i32 - %peer1, %valid1 = kernel.subgroup.shuffle %pair_sum, %xor1, %shuffle_width : f32, i32, i32 - %sum4 = scalar.addf %pair_sum, %peer1 : f32 - %peer2, %valid2 = kernel.subgroup.shuffle %sum4, %xor2, %shuffle_width : f32, i32, i32 - %sum8 = scalar.addf %sum4, %peer2 : f32 - %peer4, %valid4 = kernel.subgroup.shuffle %sum8, %xor4, %shuffle_width : f32, i32, i32 - %sum16 = scalar.addf %sum8, %peer4 : f32 - %peer8, %valid8 = kernel.subgroup.shuffle %sum16, %xor8, %shuffle_width : f32, i32, i32 - %original_wave_sum = scalar.addf %sum16, %peer8 : f32 - %original_wave = index.div %tid, %c16 : index - %cohort_lane = index.rem %tid, %c16 : index - %cohort_leader = index.cmp eq, %cohort_lane, %c0 : index - %scratch_bytes = index.constant 1024 : offset - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_view = buffer.view %scratch[%base] : buffer -> view<256xf32> - scf.if %cohort_leader { - view.store %original_wave_sum, %scratch_view[%original_wave] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %first_wave = index.cmp eq, %wave, %c0 : index - %live_partial = index.cmp ult, %lane, %c8 : index - %loads_partial = scalar.andi %first_wave, %live_partial : i1 - %partial = scf.if %loads_partial -> (f32) { - %loaded = view.load %scratch_view[%lane] : view<256xf32> -> f32 - scf.yield %loaded : f32 - } else { - scf.yield %zero : f32 - } - %row_sum = kernel.subgroup.reduce %partial : f32 - %leader = index.cmp eq, %tid, %c0 : index - scf.if %leader { - %hidden_i32 = index.cast %hidden : index to i32 - %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 - %mean = scalar.divf %row_sum, %hidden_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_view[%c0] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 - %pair_count = index.div %stripe_count, %c2 : index - %parity = index.rem %tid, %c2 : index - %odd = index.cmp ne, %parity, %c0 : index - %aligned_tid = index.sub %tid, %parity : index - %aligned_channel = index.mul %aligned_tid, %c2 : index - %stripe_lane = index.mul %parity, %c256 : index - %output_origin = index.add %aligned_channel, %stripe_lane : index - %c512 = index.constant 512 : index - scf.for %pair = [%c0 to %pair_count step %c1] unroll { - %channel0 = index.madd %pair, %c512, %output_origin : index - %channel = index.assume %channel0 [range(%channel0, 0, 8188), mul(%channel0, 4), lt(%channel0, %hidden)] : index - %stripe = index.mul %pair, %c2 : index - %phase = index.div %stripe, %c16 : index - %phase0 = index.cmp eq, %phase, %c0 : index - %phase1 = index.cmp eq, %phase, %c1 : index - %phase2 = index.cmp eq, %phase, %c2 : index - %values23 = scf.select %phase2, %cache2, %cache3 : vector<32xf32> - %values123 = scf.select %phase1, %cache1, %values23 : vector<32xf32> - %values = scf.select %phase0, %cache0, %values123 : vector<32xf32> - %cache_stripe = index.rem %stripe, %c16 : index - %cache_base = index.mul %cache_stripe, %c2 : index - %cache_index0 = index.add %cache_base, %c0 : index - %cached0 = vector.extract %values[%cache_index0] : vector<32xf32> -> f32 - %cache_index1 = index.add %cache_base, %c1 : index - %cached1 = vector.extract %values[%cache_index1] : vector<32xf32> -> f32 - %cache_index2 = index.add %cache_base, %c2 : index - %cached2 = vector.extract %values[%cache_index2] : vector<32xf32> -> f32 - %cache_index3 = index.add %cache_base, %c3 : index - %cached3 = vector.extract %values[%cache_index3] : vector<32xf32> -> f32 - %send0 = scf.select %odd, %cached0, %cached2 : f32 - %send1 = scf.select %odd, %cached1, %cached3 : f32 - %received0, %receive_valid0 = kernel.subgroup.shuffle %send0, %xor1, %shuffle_width : f32, i32, i32 - %received1, %receive_valid1 = kernel.subgroup.shuffle %send1, %xor1, %shuffle_width : f32, i32, i32 - %value0 = scf.select %odd, %received0, %cached0 : f32 - %value1 = scf.select %odd, %received1, %cached1 : f32 - %value2 = scf.select %odd, %cached2, %received0 : f32 - %value3 = scf.select %odd, %cached3, %received1 : f32 - %weights = vector.load %rhs_view[%channel] : view<[%hidden]xf32> -> vector<4xf32> - %result_zero = vector.constant 0.0 : vector<4xf32> - %weight0 = vector.extract %weights[%c0] : vector<4xf32> -> f32 - %norm0 = template.apply<@ggml.rmsnorm_f32.apply>(%value0, %row_scale) : (f32, f32) -> (f32) - %out0 = template.apply<@ggml.binary_f32.apply>(%op, %norm0, %weight0) : (index, f32, f32) -> (f32) - %result0 = vector.insert %out0 into %result_zero[%c0] : f32, vector<4xf32> - %weight1 = vector.extract %weights[%c1] : vector<4xf32> -> f32 - %norm1 = template.apply<@ggml.rmsnorm_f32.apply>(%value1, %row_scale) : (f32, f32) -> (f32) - %out1 = template.apply<@ggml.binary_f32.apply>(%op, %norm1, %weight1) : (index, f32, f32) -> (f32) - %result1 = vector.insert %out1 into %result0[%c1] : f32, vector<4xf32> - %weight2 = vector.extract %weights[%c2] : vector<4xf32> -> f32 - %norm2 = template.apply<@ggml.rmsnorm_f32.apply>(%value2, %row_scale) : (f32, f32) -> (f32) - %out2 = template.apply<@ggml.binary_f32.apply>(%op, %norm2, %weight2) : (index, f32, f32) -> (f32) - %result2 = vector.insert %out2 into %result1[%c2] : f32, vector<4xf32> - %weight3 = vector.extract %weights[%c3] : vector<4xf32> -> f32 - %norm3 = template.apply<@ggml.rmsnorm_f32.apply>(%value3, %row_scale) : (f32, f32) -> (f32) - %out3 = template.apply<@ggml.binary_f32.apply>(%op, %norm3, %weight3) : (index, f32, f32) -> (f32) - %result3 = vector.insert %out3 into %result2[%c3] : f32, vector<4xf32> - vector.store %result3, %output_view[%token, %channel] : vector<4xf32>, view<[%token_count]x[%hidden]xf32> - %packed_column = index.div %channel, %c16 : index - %packed_lane = index.rem %channel, %c16 : index - %half = vector.fptrunc %result3 : vector<4xf32> to vector<4xf16> - vector.store %half, %packed_view[%packed_column, %token, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%token_count]x16xf16> - %fence_pos = index.rem %pair, %c4 : index - %end_packet = index.cmp eq, %fence_pos, %c3 : index - scf.if %end_packet { - scf.schedule.fence - } - } - kernel.return -} +template.decl @ggml.rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_symmetric_i4_k32") @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { +template.def<@ggml.rmsnorm_binary_symmetric_i4_k32.body> device @ggml_rmsnorm_binary_symmetric_i4_k32_body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { %hidden_size0 = config.get @ggml.rmsnorm_binary_symmetric_i4.hidden_size : index %epsilon = config.get @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index @@ -436,15 +179,13 @@ kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary } } } - kernel.return + template.return } -// Publishes residual ADD, RMSNorm, learned scale, and a signed-I4/K64 alternate. -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_add_rmsnorm_binary_symmetric_i4_k32") @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { + +template.decl @ggml.add_rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.add_rmsnorm_binary_symmetric_i4_k32.body> device @ggml_add_rmsnorm_binary_symmetric_i4_k32_body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { %hidden_size0 = config.get @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : index %epsilon = config.get @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index @@ -617,15 +358,13 @@ kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_add_rmsnorm_bi } } } - kernel.return + template.return } -// Publishes RMSNorm, learned scale, SiLU gate product, and a signed-I4/K64 alternate. -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32") @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 1048576)] { + +template.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body> device @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { %hidden_size0 = config.get @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : index %epsilon = config.get @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index @@ -753,139 +492,5 @@ kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_s } } } - kernel.return -} - -template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) - -template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) - -config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] -config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 -config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] -config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] -template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) - -template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) - -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_f16") @ggml_rmsnorm_gate_f32_f16(%token_count: index) { - %one = index.constant 1 : index - %eight = index.constant 8 : index - %seven = index.constant 7 : index - %threads = index.constant 256 : index - %round = index.add %token_count, %seven : index - %groups = index.div %round, %eight : index - kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %f16_output: buffer) { - %publish_q8 = scalar.constant false : i1 - template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_q8_1_x4") @ggml_rmsnorm_gate_f32_q8_1_x4(%token_count: index) { - %one = index.constant 1 : index - %eight = index.constant 8 : index - %seven = index.constant 7 : index - %threads = index.constant 256 : index - %round = index.add %token_count, %seven : index - %groups = index.div %round, %eight : index - kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %q8_output: buffer) { - %publish_q8 = scalar.constant true : i1 - template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %q8_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -template.def<@ggml.rmsnorm_gate_f32.publish_body> device @ggml_rmsnorm_gate_f32_publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1048576)] : index - %hidden0 = config.get @ggml.rmsnorm_gate_f32.hidden_size : index - %hidden = index.assume %hidden0 [range(%hidden0, 128, 1024), mul(%hidden0, 128)] : index - %epsilon = config.get @ggml.rmsnorm_gate_f32.rms_epsilon : f32 - %gate_op = config.get @ggml.rmsnorm_gate_f32.gate_op : index - %row_width = config.get @ggml.rmsnorm_gate_f32.f16_output_row_width : index - %z = index.constant 0 : index - %zero_offset = index.constant 0 : offset - %four = index.constant 4 : index - %eight = index.constant 8 : index - %c128 = index.constant 128 : index - %wg = kernel.workgroup.id : index - %wave = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %row0 = index.madd %wg, %eight, %wave : index - %valid = index.cmp ult, %row0, %bounded_token_count : index - %safe = scf.select %valid, %row0, %z : index - %row = index.assume %safe [lt(%safe, %bounded_token_count)] : index - %channel0 = index.mul %lane, %four : index - %input_n, %weight_n, %raw_gate_n, %output_n, %alternate_n = buffer.assume.noalias %input, %weight, %raw_gate, %output, %alternate_output : buffer, buffer, buffer, buffer, buffer - %iv = buffer.view %input_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> - %wv = buffer.view %weight_n[%zero_offset] : buffer -> view<[%hidden]xf32> - %gv = buffer.view %raw_gate_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> - %ov = buffer.view %output_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> - %scale = template.apply<@ggml.rmsnorm_f32.subgroup_row_scale>(%bounded_token_count, %row, %hidden, %epsilon, %input_n) : (index, index, index, f32, buffer) -> (f32) - %scale_v = vector.splat %scale : vector<4xf32> - %last_channel = index.sub %hidden, %four : index - scf.if %valid { - scf.for %channel = [%channel0 to %hidden step %c128] { - %bounded_channel = index.assume %channel [range(%channel, 0, 1020), mul(%channel, 4), le(%channel, %last_channel)] : index - %x = vector.load %iv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> - %w = vector.load %wv[%bounded_channel] : view<[%hidden]xf32> -> vector<4xf32> - %g = vector.load %gv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> - %norm = vector.mulf %x, %scale_v : vector<4xf32> - %side = vector.mulf %norm, %w : vector<4xf32> - %activated = template.apply<@ggml.unary_f32.apply_vector4>(%gate_op, %g) : (index, vector<4xf32>) -> (vector<4xf32>) - %out = vector.mulf %side, %activated : vector<4xf32> - vector.store %out, %ov[%row, %bounded_channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden]xf32> - scf.if %publish_q8 { - %scratch_bytes = index.constant 1152 : offset - %scratch_d_offset = index.constant 1024 : offset - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%zero_offset] : buffer -> view<256xf32> - %scratch_d = buffer.view %scratch[%scratch_d_offset] : buffer -> view<32xf32> - %flat_channel = index.madd %row, %hidden, %bounded_channel : index - template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%valid, %zero_offset, %flat_channel, %out, %scratch_values, %scratch_d, %alternate_n) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) - } else { - %hv = buffer.view %alternate_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf16> - %out_half = vector.fptrunc %out : vector<4xf32> to vector<4xf16> - %packed = index.cmp ugt, %row_width, %z : index - scf.if %packed { - %c16 = index.constant 16 : index - %element_count = index.mul %bounded_token_count, %hidden : index - %packed_rows = index.div %element_count, %row_width : index - %packed_columns = index.div %row_width, %c16 : index - %packed_view = buffer.view %alternate_n[%zero_offset] : buffer -> view<[%packed_columns]x[%packed_rows]x16xf16> - %flat = index.madd %row, %hidden, %bounded_channel : index - %packed_row = index.div %flat, %row_width : index - %packed_channel = index.rem %flat, %row_width : index - %packed_column = index.div %packed_channel, %c16 : index - %packed_lane = index.rem %packed_channel, %c16 : index - vector.store %out_half, %packed_view[%packed_column, %packed_row, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%packed_rows]x16xf16> - } else { - vector.store %out_half, %hv[%row, %bounded_channel] : vector<4xf16>, view<[%bounded_token_count]x[%hidden]xf16> - } - } - } - } template.return } - -check.case public @ggml_rmsnorm_binary_f32_mul_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<128xf32> - %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> - %output = check.generate.fill value(0.0) : tensor<128xf32> - %expected = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) - check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> - check.return -} - -check.case public @ggml_rmsnorm_binary_f32_add_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<128xf32> - %rhs = check.generate.iota offset(0.0) step(0.015625) : tensor<128xf32> - %output = check.generate.fill value(0.0) : tensor<128xf32> - %expected = check.generate.iota offset(1.0) step(0.015625) : tensor<128xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) - check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> - check.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm.loom new file mode 100644 index 000000000000..f44a9d2065b0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm.loom @@ -0,0 +1,231 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) + +template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) + +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) + +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4(%arg0: i1, %arg1: offset, %arg2: index, %arg3: vector<4xf32>, %arg4: view<256xf32>, %arg5: view<32xf32>, %arg6: buffer) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) + +template.decl @ggml.rmsnorm_f32.body(%token_count: index, %input: buffer, %output: buffer) + +template.decl @ggml.rmsnorm_binary_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +template.decl @ggml.rmsnorm_binary_f32_k16.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +template.decl @ggml.rmsnorm_binary_q8_1_x4.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) + +template.decl @ggml.rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.decl @ggml.add_rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +amdgpu.target @ggml_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_f32.hidden_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] + +config.decl @ggml.rmsnorm_f32.input_stride : %value: index where [range(%value, 64, 1048576)] + +config.decl @ggml.rmsnorm_f32.rms_epsilon : f32 + +kernel.def target(@ggml_rmsnorm_gfx11_wave32) export("ggml_rmsnorm_f32") @ggml_rmsnorm_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rmsnorm_f32.body>(%token_count, %input, %output) : (index, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.fill value(1.0) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_f32[%token_count](%token_count, %input, %output) : [index](index, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + + +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32") @ggml_rmsnorm_binary_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_binary_f32.body>(%token_count, %input, %rhs, %output) : (index, buffer, buffer, buffer) + kernel.return +} + +// Retain two original lane streams per thread and publish both F32 and K16 F16. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32_k16") @ggml_rmsnorm_binary_f32_k16(%token_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_binary_f32_k16.body>(%token_count, %input, %rhs, %output, %f16_output) : (index, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_symmetric_i4_k32") @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { + template.apply<@ggml.rmsnorm_binary_symmetric_i4_k32.body>(%token_count, %input, %weight, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Publishes residual ADD, RMSNorm, learned scale, and a signed-I4/K64 alternate. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_add_rmsnorm_binary_symmetric_i4_k32") @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { + template.apply<@ggml.add_rmsnorm_binary_symmetric_i4_k32.body>(%token_count, %lhs, %rhs, %residual_output, %weight, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Publishes RMSNorm, learned scale, SiLU gate product, and a signed-I4/K64 alternate. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32") @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body>(%token_count, %input, %weight, %raw_gate, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + + +config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] +config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] +config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_f16") @ggml_rmsnorm_gate_f32_f16(%token_count: index) { + %one = index.constant 1 : index + %eight = index.constant 8 : index + %seven = index.constant 7 : index + %threads = index.constant 256 : index + %round = index.add %token_count, %seven : index + %groups = index.div %round, %eight : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %f16_output: buffer) { + %publish_q8 = scalar.constant false : i1 + template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_q8_1_x4") @ggml_rmsnorm_gate_f32_q8_1_x4(%token_count: index) { + %one = index.constant 1 : index + %eight = index.constant 8 : index + %seven = index.constant 7 : index + %threads = index.constant 256 : index + %round = index.add %token_count, %seven : index + %groups = index.div %round, %eight : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %q8_output: buffer) { + %publish_q8 = scalar.constant true : i1 + template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %q8_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + + +check.case public @ggml_rmsnorm_binary_f32_mul_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} + +check.case public @ggml_rmsnorm_binary_f32_add_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %rhs = check.generate.iota offset(0.0) step(0.015625) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(1.0) step(0.015625) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} + +amdgpu.target @ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] + + +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$28: index, %input_size$29: index) launch(%token_count$30: index, %input_size$31: index, %input: buffer, %output: buffer) + +kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnorm_binary_q8_1_x4") @ggml_rmsnorm_binary_q8_1_x4(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %hidden = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index + %tile_width = index.constant 1024 : index + %tile_round = index.constant 1023 : index + %decode_limit = index.constant 6 : index + %rounded_hidden = index.add %hidden, %tile_round : index + %stripe_count = index.div %rounded_hidden, %tile_width : index + %is_decode = index.cmp ult, %token_count, %decode_limit : index + %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index + %groups = index.mul %token_count, %groups_per_row : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rmsnorm_binary_q8_1_x4.body>(%token_count, %input, %rhs, %output, %q8_output) : (index, buffer, buffer, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_binary_q8_1_x4_mul_case { + %token_count = check.literal value(1) : index + %hidden_size = check.literal value(128) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + %f32 = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.fill value(0) : tensor<144xi8> + %actual = check.generate.fill value(1) : tensor<144xi8> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %f32) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %f32, %expected) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) + %actual_f32 = check.generate.fill value(0.0) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_q8_1_x4[%token_count](%token_count, %input, %rhs, %actual_f32, %actual) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<144xi8>) + check.expect.equal actual(%actual_f32) expected(%f32) : tensor<128xf32> + check.expect.equal actual(%actual) expected(%expected) : tensor<144xi8> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_q8_1_x4.loom deleted file mode 100644 index a4b44f79bda1..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_binary_q8_1_x4.loom +++ /dev/null @@ -1,121 +0,0 @@ -// Copyright 2026 The HRX Authors -// SPDX-License-Identifier: Apache-2.0 - -template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) - -template.decl @ggml.quantize_q8_1_x4.publish_vector4(%arg0: i1, %arg1: offset, %arg2: index, %arg3: vector<4xf32>, %arg4: view<256xf32>, %arg5: view<32xf32>, %arg6: buffer) - -template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) - -amdgpu.target @ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 - -config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] - -kernel.decl @ggml_rmsnorm_binary_f32(%token_count$23: index) launch(%token_count$24: index, %input: buffer, %rhs: buffer, %output: buffer) - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$28: index, %input_size$29: index) launch(%token_count$30: index, %input_size$31: index, %input: buffer, %output: buffer) - -kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnorm_binary_q8_1_x4") @ggml_rmsnorm_binary_q8_1_x4(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %hidden = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index - %tile_width = index.constant 1024 : index - %tile_round = index.constant 1023 : index - %decode_limit = index.constant 6 : index - %rounded_hidden = index.add %hidden, %tile_round : index - %stripe_count = index.div %rounded_hidden, %tile_width : index - %is_decode = index.cmp ult, %token_count, %decode_limit : index - %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index - %groups = index.mul %token_count, %groups_per_row : index - kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 2048)] { - %hidden_size0 = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index - %epsilon = config.get @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 - %op = config.get @ggml.rmsnorm_binary_q8_1_x4.op : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %dispatch_group = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c1024 = index.constant 1024 : index - %tile_round = index.constant 1023 : index - %decode_limit = index.constant 6 : index - %rounded_hidden = index.add %hidden_size, %tile_round : index - %stripe_count = index.div %rounded_hidden, %c1024 : index - %is_decode = index.cmp ult, %token_count, %decode_limit : index - %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index - %token0 = index.div %dispatch_group, %groups_per_row : index - %stripe_id = index.rem %dispatch_group, %groups_per_row : index - %stripe_start = index.mul %stripe_id, %c1024 : index - %stripe_step = index.mul %groups_per_row, %c1024 : index - %group_bytes = index.constant 144 : offset - %scratch_d_byte_add = index.constant 1024 : offset - %scratch_bytes = index.constant 1152 : offset - %c0_offset = index.constant 0 : offset - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %input_noalias, %rhs_noalias, %output_noalias, %q8_output_noalias = buffer.assume.noalias %input, %rhs, %output, %q8_output : buffer, buffer, buffer, buffer - %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) - %row_scale_vector = vector.splat %row_scale : vector<4xf32> - %physical_group_count = index.div %hidden_size, %c128 : index - %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset - %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> - scf.for %stripe_base = [%stripe_start to %hidden_size step %stripe_step] { - %word_element_add = index.mul %workitem, %c4 : index - %channel = index.add %stripe_base, %word_element_add : index - %valid_word0 = index.cmp ult, %channel, %hidden_size : index - %valid_word = scalar.andi %valid_word0, %valid_token : i1 - %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> - %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %rhs_values = vector.load.mask %rhs_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %normalized = vector.mulf %input_values, %row_scale_vector : vector<4xf32> - %result0 = vector.extract %normalized[0] : vector<4xf32> -> f32 - %result1 = vector.extract %normalized[1] : vector<4xf32> -> f32 - %result2 = vector.extract %normalized[2] : vector<4xf32> -> f32 - %result3 = vector.extract %normalized[3] : vector<4xf32> -> f32 - %rhs0 = vector.extract %rhs_values[0] : vector<4xf32> -> f32 - %rhs1 = vector.extract %rhs_values[1] : vector<4xf32> -> f32 - %rhs2 = vector.extract %rhs_values[2] : vector<4xf32> -> f32 - %rhs3 = vector.extract %rhs_values[3] : vector<4xf32> -> f32 - %binary0 = template.apply<@ggml.binary_f32.apply>(%op, %result0, %rhs0) : (index, f32, f32) -> (f32) - %binary1 = template.apply<@ggml.binary_f32.apply>(%op, %result1, %rhs1) : (index, f32, f32) -> (f32) - %binary2 = template.apply<@ggml.binary_f32.apply>(%op, %result2, %rhs2) : (index, f32, f32) -> (f32) - %binary3 = template.apply<@ggml.binary_f32.apply>(%op, %result3, %rhs3) : (index, f32, f32) -> (f32) - %packed0 = vector.insert %binary0 into %c0_f32x4[0] : f32, vector<4xf32> - %packed1 = vector.insert %binary1 into %packed0[1] : f32, vector<4xf32> - %packed2 = vector.insert %binary2 into %packed1[2] : f32, vector<4xf32> - %packed = vector.insert %binary3 into %packed2[3] : f32, vector<4xf32> - vector.store.mask %packed, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> - template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%valid_word, %token_output_byte_base, %channel, %packed, %scratch_values, %scratch_d, %q8_output_noalias) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) - } - kernel.return -} - -check.case public @ggml_rmsnorm_binary_q8_1_x4_mul_case { - %token_count = check.literal value(1) : index - %hidden_size = check.literal value(128) : index - %input = check.generate.fill value(2.0) : tensor<128xf32> - %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> - %f32 = check.generate.fill value(0.0) : tensor<128xf32> - %expected = check.generate.fill value(0) : tensor<144xi8> - %actual = check.generate.fill value(1) : tensor<144xi8> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %f32) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %f32, %expected) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) - %actual_f32 = check.generate.fill value(0.0) : tensor<128xf32> - kernel.launch @ggml_rmsnorm_binary_q8_1_x4[%token_count](%token_count, %input, %rhs, %actual_f32, %actual) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<144xi8>) - check.expect.equal actual(%actual_f32) expected(%f32) : tensor<128xf32> - check.expect.equal actual(%actual) expected(%expected) : tensor<144xi8> - check.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_f32.loom index 38803798b7c8..da8ba2aa5c1a 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rmsnorm_f32.loom @@ -1,57 +1,7 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) - -template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) - amdgpu.target @ggml_rmsnorm_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rmsnorm_f32.hidden_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] - -config.decl @ggml.rmsnorm_f32.input_stride : %value: index where [range(%value, 64, 1048576)] - -config.decl @ggml.rmsnorm_f32.rms_epsilon : f32 - -kernel.def target(@ggml_rmsnorm_gfx11_wave32) export("ggml_rmsnorm_f32") @ggml_rmsnorm_f32(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %hidden_size0 = config.get @ggml.rmsnorm_f32.hidden_size : index - %input_stride0 = config.get @ggml.rmsnorm_f32.input_stride : index - %epsilon = config.get @ggml.rmsnorm_f32.rms_epsilon : f32 - %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 64, 32768), mul(%hidden_size0, 64)] : index - %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 64, 1048576)] : index - %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index - %token0 = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer - %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %input_stride, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%input_stride]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - scf.if %valid_token { - scf.for %channel = [%workitem to %hidden_size step %c256] { - %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%input_stride]xf32> -> f32 - %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) - view.store %normalized, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> - } - } - kernel.return -} - -check.case public @ggml_rmsnorm_f32_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<128xf32> - %output = check.generate.fill value(0.0) : tensor<128xf32> - %expected = check.generate.fill value(1.0) : tensor<128xf32> - kernel.launch @ggml_rmsnorm_f32[%token_count](%token_count, %input, %output) : [index](index, tensor<128xf32>, tensor<128xf32>) - check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> - check.return -} // One workgroup preserves each strided RMSNorm row before scale and RoPE. config.decl @ggml.rmsnorm_mul_rope.hidden_size : %value: index where [range(%value, 2, 512)] diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json index 40e1175cfb47..99e80c399c36 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json @@ -74,15 +74,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -421,15 +421,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -474,15 +474,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -682,14 +682,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -724,14 +724,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -874,15 +874,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -1125,15 +1125,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -1300,15 +1300,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -1353,15 +1353,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -1537,14 +1537,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -1579,14 +1579,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -1729,15 +1729,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -1907,15 +1907,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -2082,15 +2082,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -2135,15 +2135,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -2319,14 +2319,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -2361,14 +2361,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -2511,15 +2511,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -2738,15 +2738,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -2913,15 +2913,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -2966,15 +2966,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -3150,14 +3150,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -3192,14 +3192,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -3342,15 +3342,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -3569,15 +3569,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -3802,15 +3802,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -3855,15 +3855,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -4039,14 +4039,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4081,14 +4081,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4231,15 +4231,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -4480,15 +4480,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -4588,15 +4588,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -4734,14 +4734,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4776,14 +4776,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4917,15 +4917,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -5424,15 +5424,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json index 558d85759a79..3425534725d4 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json @@ -74,15 +74,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -421,15 +421,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -474,15 +474,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -682,14 +682,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -724,14 +724,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -874,15 +874,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -1125,15 +1125,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -1300,15 +1300,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -1353,15 +1353,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -1537,14 +1537,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -1579,14 +1579,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -1729,15 +1729,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -1907,15 +1907,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -2082,15 +2082,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -2135,15 +2135,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -2319,14 +2319,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -2361,14 +2361,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -2511,15 +2511,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -2738,15 +2738,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -2913,15 +2913,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -2966,15 +2966,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -3150,14 +3150,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -3192,14 +3192,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -3342,15 +3342,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -3569,15 +3569,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -3802,15 +3802,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -3855,15 +3855,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -4039,14 +4039,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4081,14 +4081,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4231,15 +4231,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -4480,15 +4480,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", @@ -4588,15 +4588,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -4734,14 +4734,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4776,14 +4776,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -4917,15 +4917,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -5444,15 +5444,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json index 57e4131da09b..521002c7d7ee 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json @@ -74,15 +74,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", @@ -367,15 +367,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", @@ -420,15 +420,15 @@ }, "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", @@ -604,14 +604,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -646,14 +646,14 @@ "integer_parameters": {}, "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "primary_sources": [ "ops/rmsnorm_f32.loom" ], "sources": [ "ops/rmsnorm_f32.loom", - "motifs/rmsnorm_f32.loom" + "motifs/normalization/rmsnorm.loom" ], "symbol": "ggml_rmsnorm_mul_rope_f32", "workload_parameters": [] @@ -796,15 +796,15 @@ }, "kernel": "loom_libs:ggml_rmsnorm_binary_f32", "library_sources": [ - "motifs/rmsnorm_f32.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "primary_sources": [ - "ops/rmsnorm_binary_f32.loom" + "ops/normalization/rmsnorm.loom" ], "sources": [ - "ops/rmsnorm_binary_f32.loom", - "motifs/rmsnorm_f32.loom", + "ops/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm.loom", "motifs/binary_f32_apply.loom" ], "symbol": "ggml_rmsnorm_binary_f32", From 8193e467505f79373108b848dcbf0860c41211de Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 11:34:35 -0700 Subject: [PATCH 4/8] Clean up HRX ROPE and rows Loom kernels Group ROPE and rows Loom sources under dedicated op and motif directories, fold the fused ROPE+set_rows kernel into the ROPE source, and share cache row publication helpers between set_rows and the fused path. Update the HRX Loom manifest and checked-in benchmark program JSONs for the new source layout. Assisted-by: Codex --- .../benchmarks/loom/llama32_3b_f16.pp256.json | 12 +- .../benchmarks/loom/llama32_3b_f16.pp512.json | 12 +- .../benchmarks/loom/llama32_3b_f16.tg8.json | 33 +- .../loom/qwen3_30b_a3b_q4_k_m.pp256.json | 4 +- .../loom/qwen3_30b_a3b_q4_k_m.pp512.json | 4 +- .../loom/qwen3_30b_a3b_q4_k_m.tg8.json | 4 +- .../common/dispatch-rope-set-rows.cpp | 57 ++-- .../kernels/loom-libs/manifest.json | 78 +++-- .../loom-libs/motifs/{ => rope}/rope_f32.loom | 0 .../loom-libs/motifs/rows/cache_publish.loom | 69 ++++ .../kernels/loom-libs/ops/rope/rope_f32.loom | 306 ++++++++++++++++++ .../kernels/loom-libs/ops/rope_f32.loom | 150 --------- .../loom-libs/ops/rope_set_rows_f32.loom | 209 ------------ .../ops/{ => rows}/get_rows_f32.loom | 0 .../loom-libs/ops/{ => rows}/set_rows.loom | 43 +-- .../loom/v2_draft_programs.tg_c1.json | 114 ++++--- .../loom/v2_mtp4_companion_all.tg_c5.json | 114 ++++--- .../loom/v2_mtp4_program4.tg_c4.json | 25 +- 18 files changed, 631 insertions(+), 603 deletions(-) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/{ => rope}/rope_f32.loom (100%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_set_rows_f32.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{ => rows}/get_rows_f32.loom (100%) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{ => rows}/set_rows.loom (70%) diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json index 1b7175649af0..18306b93d7e2 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json @@ -29,10 +29,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -111,7 +111,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -128,7 +128,7 @@ "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom", "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -219,7 +219,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -236,7 +236,7 @@ "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom", "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json index 930df3951246..9752830da780 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json @@ -29,10 +29,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -111,7 +111,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -128,7 +128,7 @@ "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom", "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -219,7 +219,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -236,7 +236,7 @@ "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom", "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json index 96c3b274f7e7..6aa8f92f4498 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json @@ -29,10 +29,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -213,14 +213,16 @@ }, "kernel": "loom_libs:ggml_rope_f32", "library_sources": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ], "primary_sources": [ - "ops/rope_f32.loom" + "ops/rope/rope_f32.loom" ], "sources": [ - "ops/rope_f32.loom", - "motifs/rope_f32.loom" + "ops/rope/rope_f32.loom", + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_rope_f32", "workload_parameters": [ @@ -247,14 +249,16 @@ }, "kernel": "loom_libs:ggml_rope_set_rows_f32", "library_sources": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ], "primary_sources": [ - "ops/rope_set_rows_f32.loom" + "ops/rope/rope_f32.loom" ], "sources": [ - "ops/rope_set_rows_f32.loom", - "motifs/rope_f32.loom" + "ops/rope/rope_f32.loom", + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_rope_set_rows_f32", "workload_parameters": [ @@ -284,12 +288,15 @@ "token_count": 1 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json index b0f94c24680a..912666842ec6 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json @@ -77,10 +77,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json index f3734c0b1410..d5eaa64c859b 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json @@ -77,10 +77,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json index 35cbb8e814a9..77793a82d7af 100644 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json @@ -75,10 +75,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp index 860feb88034a..0399fe094524 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp @@ -265,12 +265,12 @@ struct SetRowsMatch { struct RopeSetRowsMatch { RopeMatch rope; SetRowsMatch set_rows; - const GraphNode * layout = nullptr; - const Value * cache_rows = nullptr; - size_t set_rows_idx = 0; - size_t layout_idx = 0; + const GraphNode * layout = nullptr; + const Value * published_rows = nullptr; + size_t set_rows_idx = 0; + size_t layout_idx = 0; - bool matched() const { return rope.matched() && set_rows.matched() && cache_rows != nullptr; } + bool matched() const { return rope.matched() && set_rows.matched() && published_rows != nullptr; } }; static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Status * status = nullptr) { @@ -417,15 +417,20 @@ static SetRowsMatch match_set_rows_2d(const Graph & graph, const GraphNode * nod return match; } -static void add_rope_compile_parameters(Dispatch & dispatch, const RopeMatch & match) { - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.head_size", to_config_value(match.head_size)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.n_dims", to_config_value(match.n_dims)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.head_count", to_config_value(match.head_count)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.token_capacity", to_config_value(match.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.input_stride1", to_config_value(match.input_stride1)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.input_stride2", to_config_value(match.input_stride2)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.mscale", rope_mscale_config_value(match.rope_mscale)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.mode", to_config_value(match.mode)); +static void add_rope_compile_parameters(Dispatch & dispatch, const RopeMatch & match, const char * prefix) { + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_size", to_config_value(match.head_size)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".n_dims", to_config_value(match.n_dims)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_count", to_config_value(match.head_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".token_capacity", + to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_stride1", + to_config_value(match.input_stride1)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_stride2", + to_config_value(match.input_stride2)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_span", to_config_value(match.input_span)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".mscale", + rope_mscale_config_value(match.rope_mscale)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".mode", to_config_value(match.mode)); } static void add_set_rows_compile_parameters(Dispatch & dispatch, const SetRowsMatch & match) { @@ -476,8 +481,7 @@ static Dispatch make_rope_dispatch(const DispatchMatchContext & context, Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRopeF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - add_rope_compile_parameters(dispatch, match); - dispatch.kernel.compile_parameters.emplace("ggml.rope_f32.input_span", to_config_value(match.input_span)); + add_rope_compile_parameters(dispatch, match, "ggml.rope_f32"); dispatch.bindings.push_back({ match.positions->id, 0, match.positions->byte_count }); dispatch.bindings.push_back({ match.input->storage_root, match.input->storage_offset, match.input_span_bytes }); dispatch.bindings.push_back({ theta, 0, match.theta_bytes }); @@ -508,24 +512,9 @@ static Dispatch make_rope_set_rows_dispatch(const DispatchMatchContext & context dispatch.kernel = make_kernel_specialization(kRopeSetRowsF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); dispatch.kernel.integer_parameters.emplace("cache_row_count", match.set_rows.cache_row_count); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.head_size", - to_config_value(match.rope.head_size)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.n_dims", to_config_value(match.rope.n_dims)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.head_count", - to_config_value(match.rope.head_count)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.token_capacity", - to_config_value(match.rope.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.input_stride1", - to_config_value(match.rope.input_stride1)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.input_stride2", - to_config_value(match.rope.input_stride2)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.input_span", - to_config_value(match.rope.input_span)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.mscale", - rope_mscale_config_value(match.rope.rope_mscale)); + add_rope_compile_parameters(dispatch, match.rope, "ggml.rope_set_rows_f32"); dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.output_format", to_config_value(match.set_rows.output_format)); - dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.mode", to_config_value(match.rope.mode)); dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); dispatch.bindings.push_back({ match.set_rows.indices->id, 0, match.set_rows.indices->byte_count }); dispatch.bindings.push_back( @@ -574,9 +563,9 @@ static RopeSetRowsMatch match_rope_set_rows_f32(const Graph & graph, const Graph return {}; } - match.cache_rows = match.set_rows.rows; + match.published_rows = match.set_rows.rows; if (match.set_rows.row_format != 32 || match.set_rows.hidden_size != match.rope.head_size * match.rope.head_count || - match.set_rows.token_count != match.rope.token_count || match.cache_rows->type != GGML_TYPE_F32) { + match.set_rows.token_count != match.rope.token_count || match.published_rows->type != GGML_TYPE_F32) { return {}; } return match; diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 53528b3f251f..c5c468de228a 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -147,7 +147,7 @@ "path": "ops/mul_mat_f32_f32_decode.loom" }, { - "path": "ops/get_rows_f32.loom" + "path": "ops/rows/get_rows_f32.loom" }, { "path": "ops/gated_delta_net_f32_wmma.loom" @@ -159,13 +159,10 @@ "path": "ops/ssm_conv_generic_f32.loom" }, { - "path": "ops/rope_f32.loom" + "path": "ops/rope/rope_f32.loom" }, { - "path": "ops/set_rows.loom" - }, - { - "path": "ops/rope_set_rows_f32.loom" + "path": "ops/rows/set_rows.loom" }, { "path": "ops/flash_attention_f32_f16_wmma.loom" @@ -198,7 +195,10 @@ "path": "motifs/publish_f32.loom" }, { - "path": "motifs/rope_f32.loom" + "path": "motifs/rope/rope_f32.loom" + }, + { + "path": "motifs/rows/cache_publish.loom" }, { "path": "motifs/dequant.loom" @@ -1993,7 +1993,7 @@ "library_sources": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2007,7 +2007,7 @@ "compile_dependencies": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2062,7 +2062,7 @@ "library_sources": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2076,7 +2076,7 @@ "compile_dependencies": [ "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2190,7 +2190,7 @@ "library_sources": [ "motifs/matmul/vector_activation.loom", "motifs/matmul/vector_decode_pair.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/matmul/vector_publish.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", @@ -2206,7 +2206,7 @@ "compile_dependencies": [ "motifs/matmul/vector_activation.loom", "motifs/matmul/vector_decode_pair.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/matmul/vector_publish.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", @@ -2263,7 +2263,7 @@ "library_sources": [ "motifs/matmul/vector_activation.loom", "motifs/matmul/vector_decode_pair.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/matmul/vector_publish.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", @@ -2279,7 +2279,7 @@ "compile_dependencies": [ "motifs/matmul/vector_activation.loom", "motifs/matmul/vector_decode_pair.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/matmul/vector_publish.loom", "motifs/llm_attention_qkv_matmul_postops.loom", "motifs/dequant.loom", @@ -2399,7 +2399,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -2414,7 +2414,7 @@ "compile_dependencies": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -2470,7 +2470,7 @@ "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -2485,7 +2485,7 @@ "compile_dependencies": [ "motifs/mul_mat_f32_f32_wmma_core.loom", "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope_f32.loom", + "motifs/rope/rope_f32.loom", "motifs/dequant.loom", "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", @@ -5265,7 +5265,7 @@ "name": "ggml_get_rows_f32", "family": "loom_libs", "symbol": "ggml_get_rows_f32", - "source": "ops/get_rows_f32.loom", + "source": "ops/rows/get_rows_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -5308,7 +5308,7 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "library_sources": [ "motifs/dequant.loom", @@ -5340,7 +5340,7 @@ "name": "ggml_get_rows_f32_next", "family": "loom_libs", "symbol": "ggml_get_rows_f32_next", - "source": "ops/get_rows_f32.loom", + "source": "ops/rows/get_rows_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -5385,7 +5385,7 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "library_sources": [ "motifs/dequant.loom", @@ -5417,7 +5417,7 @@ "name": "ggml_rope_f32", "family": "loom_libs", "symbol": "ggml_rope_f32", - "source": "ops/rope_f32.loom", + "source": "ops/rope/rope_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -5448,21 +5448,23 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rope_f32.loom" + "ops/rope/rope_f32.loom" ], "library_sources": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ] }, "compile_dependencies": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ] }, { "name": "ggml_set_rows", "family": "loom_libs", "symbol": "ggml_set_rows", - "source": "ops/set_rows.loom", + "source": "ops/rows/set_rows.loom", "workload_parameters": [ { "name": "token_count", @@ -5505,17 +5507,21 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], - "library_sources": [] + "library_sources": [ + "motifs/rows/cache_publish.loom" + ] }, - "compile_dependencies": [] + "compile_dependencies": [ + "motifs/rows/cache_publish.loom" + ] }, { "name": "ggml_rope_set_rows_f32", "family": "loom_libs", "symbol": "ggml_rope_set_rows_f32", - "source": "ops/rope_set_rows_f32.loom", + "source": "ops/rope/rope_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -5556,14 +5562,16 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/rope_set_rows_f32.loom" + "ops/rope/rope_f32.loom" ], "library_sources": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ] }, "compile_dependencies": [ - "motifs/rope_f32.loom" + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" ] }, { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom similarity index 100% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope_f32.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom new file mode 100644 index 000000000000..59eaa9d8c006 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rows.cache_publish.safe_i64_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector2(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<2xf32>, %cache: buffer) + +template.decl @ggml.rows.cache_publish.f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) + +template.def<@ggml.rows.cache_publish.safe_i64_row> device @ggml_rows_cache_publish_safe_i64_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %c0 = index.constant 0 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %indices_view = buffer.view %indices[%c0_offset] : buffer -> view<[%bounded_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%bounded_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %bounded_cache_row_count : index + %publish = scalar.andi %valid_index, %valid_row : i1 + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + template.return %publish, %cache_row : i1, index +} + +template.def<@ggml.rows.cache_publish.f32_vector2> device @ggml_rows_cache_publish_f32_vector2(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<2xf32>, %cache: buffer) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %output_format, %c16 : index + %is_f32 = index.cmp eq, %output_format, %c32 : index + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf16> + %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<2xf16>, view<[%cache_row_count]x[%hidden_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<2xf32>, view<[%cache_row_count]x[%hidden_size]xf32> + } + template.return +} + +template.def<@ggml.rows.cache_publish.f32_vector4> device @ggml_rows_cache_publish_f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %output_format, %c16 : index + %is_f32 = index.cmp eq, %output_format, %c32 : index + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf16> + %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<4xf16>, view<[%cache_row_count]x[%hidden_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<4xf32>, view<[%cache_row_count]x[%hidden_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom new file mode 100644 index 000000000000..69e1bb57ad34 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom @@ -0,0 +1,306 @@ +// Generic GGML ROPE for F32 head-major rows. +template.decl @ggml.rope_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) + +template.decl @ggml.rope_set_rows_f32.body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) + +template.decl @ggml.rope_f32.rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +template.decl @ggml.rope_f32.rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) + +template.decl @ggml.rope_f32.copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) + +template.decl @ggml.rows.cache_publish.safe_i64_row(%arg0: index, %arg1: index, %arg2: index, %arg3: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector2(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<2xf32>, %arg7: buffer) + +template.decl @ggml.rows.cache_publish.f32_vector4(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) + +amdgpu.target @ggml_rope_f32_gfx11_wave32 {subgroup_size = 32} + +amdgpu.target @ggml_rope_set_rows_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.rope_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_f32.input_span : %value: index where [range(%value, 1, 1073741824)] + +config.decl @ggml.rope_f32.mscale : f32 + +config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] + +config.decl @ggml.rope_set_rows_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_set_rows_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_set_rows_f32.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.rope_set_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.rope_set_rows_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.input_span : %value: index where [range(%value, 1, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.mscale : f32 + +config.decl @ggml.rope_set_rows_f32.output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.rope_set_rows_f32.mode : %value: index where [range(%value, 0, 2)] + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +template.def<@ggml.rope_f32.rotate_neox_packet> device @ggml_rope_f32_rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %c2 = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %mscale_vector = vector.splat %mscale : vector<2xf32> + %paired_channel = index.add %pair_channel, %half_n_dims : index + %low_index0 = index.add %input_base, %pair_channel : index + %high_index0 = index.add %input_base, %paired_channel : index + %low_end0 = index.add %low_index0, %c2 : index + %high_end0 = index.add %high_index0, %c2 : index + %low_index, %low_end = index.assume %low_index0, %low_end0 [le(%low_end0, %bounded_input_span)] : index, index + %high_index, %high_end = index.assume %high_index0, %high_end0 [le(%high_end0, %bounded_input_span)] : index, index + %low_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%low_end]xf32> + %high_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%high_end]xf32> + %low_values = vector.load %low_input_view[%low_index] : view<[%low_end]xf32> -> vector<2xf32> + %high_values = vector.load %high_input_view[%high_index] : view<[%high_end]xf32> -> vector<2xf32> + %rotated_low, %rotated_high = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %low_values, %high_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %scaled_low = vector.mulf %rotated_low, %mscale_vector : vector<2xf32> + %scaled_high = vector.mulf %rotated_high, %mscale_vector : vector<2xf32> + template.return %scaled_low, %scaled_high : vector<2xf32>, vector<2xf32> +} + +template.def<@ggml.rope_f32.rotate_normal_packet> device @ggml_rope_f32_rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) { + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %mscale_vector = vector.splat %mscale : vector<2xf32> + %normal_index0 = index.add %input_base, %normal_channel : index + %normal_end0 = index.add %normal_index0, %c4 : index + %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index + %normal_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%normal_end]xf32> + %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> + %x0 = vector.extract %values[0] : vector<4xf32> -> f32 + %y0 = vector.extract %values[1] : vector<4xf32> -> f32 + %x1 = vector.extract %values[2] : vector<4xf32> -> f32 + %y1 = vector.extract %values[3] : vector<4xf32> -> f32 + %x_values = vector.from_elements %x0, %x1 : vector<2xf32> + %y_values = vector.from_elements %y0, %y1 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %scaled_x = vector.mulf %rotated_x, %mscale_vector : vector<2xf32> + %scaled_y = vector.mulf %rotated_y, %mscale_vector : vector<2xf32> + %scaled_x0 = vector.extract %scaled_x[0] : vector<2xf32> -> f32 + %scaled_x1 = vector.extract %scaled_x[1] : vector<2xf32> -> f32 + %scaled_y0 = vector.extract %scaled_y[0] : vector<2xf32> -> f32 + %scaled_y1 = vector.extract %scaled_y[1] : vector<2xf32> -> f32 + %rotated = vector.from_elements %scaled_x0, %scaled_y0, %scaled_x1, %scaled_y1 : vector<4xf32> + template.return %rotated : vector<4xf32> +} + +template.def<@ggml.rope_f32.copy_packet> device @ggml_rope_f32_copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) { + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %normal_index0 = index.add %input_base, %normal_channel : index + %normal_end0 = index.add %normal_index0, %c4 : index + %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index + %normal_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%normal_end]xf32> + %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> + template.return %values : vector<4xf32> +} + +template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count0 = config.get @ggml.rope_f32.head_count : index + %head_size0 = config.get @ggml.rope_f32.head_size : index + %n_dims0 = config.get @ggml.rope_f32.n_dims : index + %input_stride1_0 = config.get @ggml.rope_f32.input_stride1 : index + %input_stride2_0 = config.get @ggml.rope_f32.input_stride2 : index + %input_span0 = config.get @ggml.rope_f32.input_span : index + %mscale = config.get @ggml.rope_f32.mscale : f32 + %mode = config.get @ggml.rope_f32.mode : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index + %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index + %head = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %half_n_dims = index.div %n_dims, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index + %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index + %copy_channel = scalar.andi %active_channel, %tail_channel : i1 + %publish = scalar.andi %valid_token, %active_channel : i1 + %publish_rotary0 = scalar.andi %valid_token, %rotary_channel : i1 + %publish_rotary = scalar.andi %publish_rotary0, %active_channel : i1 + %publish_copy = scalar.andi %valid_token, %copy_channel : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %is_normal = index.cmp eq, %mode, %c0 : index + %positions_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %input, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer, buffer + scf.if %publish_rotary { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %normal_channel = index.mul %channel, %c4 : index + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + vector.store %scaled_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + vector.store %scaled_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + scf.if %is_normal { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %input_base, %normal_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + vector.store %rotated, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + } + scf.if %publish_copy { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %normal_channel = index.mul %channel, %c4 : index + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %normal_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + template.return +} + +kernel.def target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count = config.get @ggml.rope_f32.head_count : index + %head_size = config.get @ggml.rope_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %output) : (index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +template.def<@ggml.rope_set_rows_f32.body> device @ggml_rope_set_rows_f32_body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index + %head_count0 = config.get @ggml.rope_set_rows_f32.head_count : index + %head_size0 = config.get @ggml.rope_set_rows_f32.head_size : index + %n_dims0 = config.get @ggml.rope_set_rows_f32.n_dims : index + %input_stride1_0 = config.get @ggml.rope_set_rows_f32.input_stride1 : index + %input_stride2_0 = config.get @ggml.rope_set_rows_f32.input_stride2 : index + %input_span0 = config.get @ggml.rope_set_rows_f32.input_span : index + %mscale = config.get @ggml.rope_set_rows_f32.mscale : f32 + %output_format = config.get @ggml.rope_set_rows_f32.output_format : index + %mode = config.get @ggml.rope_set_rows_f32.mode : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index + %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index + %head = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %half_n_dims = index.div %n_dims, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %cache_hidden_size = index.mul %head_count, %head_size : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index + %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index + %copy_channel = scalar.andi %active_channel, %tail_channel : i1 + %publish0 = scalar.andi %valid_token, %active_channel : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %is_normal = index.cmp eq, %mode, %c0 : index + %positions_noalias, %indices_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %input, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %valid_row, %cache_row = template.apply<@ggml.rows.cache_publish.safe_i64_row>(%token, %bounded_token_count, %bounded_cache_row_count, %indices_noalias) : (index, index, index, buffer) -> (i1, index) + %publish = scalar.andi %publish0, %valid_row : i1 + %publish_rotary = scalar.andi %publish, %rotary_channel : i1 + %publish_copy = scalar.andi %publish, %copy_channel : i1 + scf.if %publish_rotary { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %normal_channel = index.mul %channel, %c4 : index + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %cache_head_base = index.mul %head, %head_size : index + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %low_cache_channel = index.add %cache_head_base, %pair_channel : index + %high_cache_channel = index.add %cache_head_base, %paired_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector2>(%output_format, %publish, %cache_row, %low_cache_channel, %bounded_cache_row_count, %cache_hidden_size, %scaled_low, %cache_noalias) : (index, i1, index, index, index, index, vector<2xf32>, buffer) + template.apply<@ggml.rows.cache_publish.f32_vector2>(%output_format, %publish, %cache_row, %high_cache_channel, %bounded_cache_row_count, %cache_hidden_size, %scaled_high, %cache_noalias) : (index, i1, index, index, index, index, vector<2xf32>, buffer) + } + scf.if %is_normal { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %input_base, %normal_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + %cache_channel = index.add %cache_head_base, %normal_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %cache_channel, %bounded_cache_row_count, %cache_hidden_size, %rotated, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + } + } + scf.if %publish_copy { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %normal_channel = index.mul %channel, %c4 : index + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %normal_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + %cache_head_base = index.mul %head, %head_size : index + %cache_channel = index.add %cache_head_base, %normal_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %cache_channel, %bounded_cache_row_count, %cache_hidden_size, %values, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +kernel.def target(@ggml_rope_set_rows_f32_gfx11_wave32) @ggml_rope_set_rows_f32(%token_count: index, %cache_row_count: index) { + %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index + %head_count = config.get @ggml.rope_set_rows_f32.head_count : index + %head_size = config.get @ggml.rope_set_rows_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_set_rows_f32.body>(%token_count, %cache_row_count, %positions, %indices, %input, %theta, %freq_factors, %cache) : (index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_f32.loom deleted file mode 100644 index 78d54b38e7d1..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_f32.loom +++ /dev/null @@ -1,150 +0,0 @@ -// Generic GGML ROPE for F32 head-major rows. -template.decl @ggml.rope_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) - -amdgpu.target @ggml_rope_f32_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] - -config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] - -config.decl @ggml.rope_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] - -config.decl @ggml.rope_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] - -config.decl @ggml.rope_f32.input_span : %value: index where [range(%value, 1, 1073741824)] - -config.decl @ggml.rope_f32.mscale : f32 - -config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] - -func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - -template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { - %token_capacity = config.get @ggml.rope_f32.token_capacity : index - %head_count0 = config.get @ggml.rope_f32.head_count : index - %head_size0 = config.get @ggml.rope_f32.head_size : index - %n_dims0 = config.get @ggml.rope_f32.n_dims : index - %input_stride1_0 = config.get @ggml.rope_f32.input_stride1 : index - %input_stride2_0 = config.get @ggml.rope_f32.input_stride2 : index - %input_span0 = config.get @ggml.rope_f32.input_span : index - %mscale = config.get @ggml.rope_f32.mscale : f32 - %mode = config.get @ggml.rope_f32.mode : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index - %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index - %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index - %head = kernel.workgroup.id : index - %token0 = kernel.workgroup.id : index - %channel0 = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %mode_neox = index.constant 2 : index - %c0_offset = index.constant 0 : offset - %half_n_dims = index.div %n_dims, %c2 : index - %pair_packet_count = index.div %head_size, %c4 : index - %rotary_packet_count = index.div %n_dims, %c4 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %active_channel = index.cmp ult, %channel0, %pair_packet_count : index - %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index - %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index - %copy_channel = scalar.andi %active_channel, %tail_channel : i1 - %publish = scalar.andi %valid_token, %active_channel : i1 - %publish_rotary0 = scalar.andi %valid_token, %rotary_channel : i1 - %publish_rotary = scalar.andi %publish_rotary0, %active_channel : i1 - %publish_copy = scalar.andi %valid_token, %copy_channel : i1 - %token = scf.select %valid_token, %token0, %c0 : index - %is_neox = index.cmp eq, %mode, %mode_neox : index - %is_normal = index.cmp eq, %mode, %c0 : index - %positions_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %input, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer, buffer - scf.if %publish_rotary { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %pair_channel = index.mul %channel, %c2 : index - %normal_channel = index.mul %channel, %c4 : index - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> - %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> - %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - %input_token_offset = index.mul %token, %input_stride2 : index - %input_head_offset = index.mul %head, %input_stride1 : index - %input_base = index.add %input_token_offset, %input_head_offset : index - %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> - %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> - scf.if %is_neox { - %mscale_vector = vector.splat %mscale : vector<2xf32> - %paired_channel = index.add %pair_channel, %half_n_dims : index - %low_index0 = index.add %input_base, %pair_channel : index - %high_index0 = index.add %input_base, %paired_channel : index - %low_end0 = index.add %low_index0, %c2 : index - %high_end0 = index.add %high_index0, %c2 : index - %low_index, %low_end = index.assume %low_index0, %low_end0 [le(%low_end0, %bounded_input_span)] : index, index - %high_index, %high_end = index.assume %high_index0, %high_end0 [le(%high_end0, %bounded_input_span)] : index, index - %low_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%low_end]xf32> - %high_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%high_end]xf32> - %low_values = vector.load %low_input_view[%low_index] : view<[%low_end]xf32> -> vector<2xf32> - %high_values = vector.load %high_input_view[%high_index] : view<[%high_end]xf32> -> vector<2xf32> - %rotated_low, %rotated_high = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %low_values, %high_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %scaled_low = vector.mulf %rotated_low, %mscale_vector : vector<2xf32> - %scaled_high = vector.mulf %rotated_high, %mscale_vector : vector<2xf32> - vector.store %scaled_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - vector.store %scaled_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - } - scf.if %is_normal { - %mscale_vector = vector.splat %mscale : vector<2xf32> - %normal_index0 = index.add %input_base, %normal_channel : index - %normal_end0 = index.add %normal_index0, %c4 : index - %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index - %normal_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%normal_end]xf32> - %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> - %x0 = vector.extract %values[0] : vector<4xf32> -> f32 - %y0 = vector.extract %values[1] : vector<4xf32> -> f32 - %x1 = vector.extract %values[2] : vector<4xf32> -> f32 - %y1 = vector.extract %values[3] : vector<4xf32> -> f32 - %x_values = vector.from_elements %x0, %x1 : vector<2xf32> - %y_values = vector.from_elements %y0, %y1 : vector<2xf32> - %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %scaled_x = vector.mulf %rotated_x, %mscale_vector : vector<2xf32> - %scaled_y = vector.mulf %rotated_y, %mscale_vector : vector<2xf32> - %scaled_x0 = vector.extract %scaled_x[0] : vector<2xf32> -> f32 - %scaled_x1 = vector.extract %scaled_x[1] : vector<2xf32> -> f32 - %scaled_y0 = vector.extract %scaled_y[0] : vector<2xf32> -> f32 - %scaled_y1 = vector.extract %scaled_y[1] : vector<2xf32> -> f32 - %rotated = vector.from_elements %scaled_x0, %scaled_y0, %scaled_x1, %scaled_y1 : vector<4xf32> - vector.store %rotated, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - } - } - scf.if %publish_copy { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %normal_channel = index.mul %channel, %c4 : index - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - %input_token_offset = index.mul %token, %input_stride2 : index - %input_head_offset = index.mul %head, %input_stride1 : index - %input_base = index.add %input_token_offset, %input_head_offset : index - %normal_index0 = index.add %input_base, %normal_channel : index - %normal_end0 = index.add %normal_index0, %c4 : index - %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index - %normal_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%normal_end]xf32> - %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> - vector.store %values, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> - } - template.return -} - -kernel.def target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count: index) { - %token_capacity = config.get @ggml.rope_f32.token_capacity : index - %head_count = config.get @ggml.rope_f32.head_count : index - %head_size = config.get @ggml.rope_f32.head_size : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %pair_packet_count = index.div %head_size, %c4 : index - kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index -} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - template.apply<@ggml.rope_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %output) : (index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_set_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_set_rows_f32.loom deleted file mode 100644 index fb649c0f5361..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope_set_rows_f32.loom +++ /dev/null @@ -1,209 +0,0 @@ -// Fused F32 ROPE followed by indexed cache publication. -template.decl @ggml.rope_set_rows_f32.body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) - -template.decl @ggml.rope_set_rows_f32.store(%output_format: index, %publish: i1, %cache_row: index, %head: index, %channel: index, %cache_row_count: index, %head_count: index, %head_size: index, %values: vector<2xf32>, %cache: buffer) - -template.decl @ggml.rope_set_rows_f32.store4(%output_format: index, %publish: i1, %cache_row: index, %head: index, %channel: index, %cache_row_count: index, %head_count: index, %head_size: index, %values: vector<4xf32>, %cache: buffer) - -amdgpu.target @ggml_rope_set_rows_f32_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rope_set_rows_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @ggml.rope_set_rows_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] - -config.decl @ggml.rope_set_rows_f32.head_count : %value: index where [range(%value, 1, 64)] - -config.decl @ggml.rope_set_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] - -config.decl @ggml.rope_set_rows_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] - -config.decl @ggml.rope_set_rows_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] - -config.decl @ggml.rope_set_rows_f32.input_span : %value: index where [range(%value, 1, 1073741824)] - -config.decl @ggml.rope_set_rows_f32.mscale : f32 - -config.decl @ggml.rope_set_rows_f32.output_format : %value: index where [range(%value, 16, 32)] - -config.decl @ggml.rope_set_rows_f32.mode : %value: index where [range(%value, 0, 2)] - -func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - -template.def<@ggml.rope_set_rows_f32.store> device @ggml_rope_set_rows_f32_store(%output_format: index, %publish: i1, %cache_row: index, %head: index, %channel: index, %cache_row_count: index, %head_count: index, %head_size: index, %values: vector<2xf32>, %cache: buffer) { - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c0_offset = index.constant 0 : offset - %is_f16 = index.cmp eq, %output_format, %c16 : index - %is_f32 = index.cmp eq, %output_format, %c32 : index - %publish_f16 = scalar.andi %publish, %is_f16 : i1 - %publish_f32 = scalar.andi %publish, %is_f32 : i1 - scf.if %publish_f16 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%head_count]x[%head_size]xf16> - %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> - vector.store %truncated, %cache_view[%cache_row, %head, %channel] : vector<2xf16>, view<[%cache_row_count]x[%head_count]x[%head_size]xf16> - } - scf.if %publish_f32 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%head_count]x[%head_size]xf32> - vector.store %values, %cache_view[%cache_row, %head, %channel] : vector<2xf32>, view<[%cache_row_count]x[%head_count]x[%head_size]xf32> - } - template.return -} - -template.def<@ggml.rope_set_rows_f32.store4> device @ggml_rope_set_rows_f32_store4(%output_format: index, %publish: i1, %cache_row: index, %head: index, %channel: index, %cache_row_count: index, %head_count: index, %head_size: index, %values: vector<4xf32>, %cache: buffer) { - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c0_offset = index.constant 0 : offset - %is_f16 = index.cmp eq, %output_format, %c16 : index - %is_f32 = index.cmp eq, %output_format, %c32 : index - %publish_f16 = scalar.andi %publish, %is_f16 : i1 - %publish_f32 = scalar.andi %publish, %is_f32 : i1 - scf.if %publish_f16 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%head_count]x[%head_size]xf16> - %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> - vector.store %truncated, %cache_view[%cache_row, %head, %channel] : vector<4xf16>, view<[%cache_row_count]x[%head_count]x[%head_size]xf16> - } - scf.if %publish_f32 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%head_count]x[%head_size]xf32> - vector.store %values, %cache_view[%cache_row, %head, %channel] : vector<4xf32>, view<[%cache_row_count]x[%head_count]x[%head_size]xf32> - } - template.return -} - -template.def<@ggml.rope_set_rows_f32.body> device @ggml_rope_set_rows_f32_body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { - %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index - %head_count0 = config.get @ggml.rope_set_rows_f32.head_count : index - %head_size0 = config.get @ggml.rope_set_rows_f32.head_size : index - %n_dims0 = config.get @ggml.rope_set_rows_f32.n_dims : index - %input_stride1_0 = config.get @ggml.rope_set_rows_f32.input_stride1 : index - %input_stride2_0 = config.get @ggml.rope_set_rows_f32.input_stride2 : index - %input_span0 = config.get @ggml.rope_set_rows_f32.input_span : index - %mscale = config.get @ggml.rope_set_rows_f32.mscale : f32 - %output_format = config.get @ggml.rope_set_rows_f32.output_format : index - %mode = config.get @ggml.rope_set_rows_f32.mode : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index - %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index - %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index - %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index - %head = kernel.workgroup.id : index - %token0 = kernel.workgroup.id : index - %channel0 = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %mode_neox = index.constant 2 : index - %c0_i64 = scalar.constant 0 : i64 - %c1048575_i64 = scalar.constant 1048575 : i64 - %c0_offset = index.constant 0 : offset - %half_n_dims = index.div %n_dims, %c2 : index - %pair_packet_count = index.div %head_size, %c4 : index - %rotary_packet_count = index.div %n_dims, %c4 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %active_channel = index.cmp ult, %channel0, %pair_packet_count : index - %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index - %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index - %copy_channel = scalar.andi %active_channel, %tail_channel : i1 - %publish0 = scalar.andi %valid_token, %active_channel : i1 - %token = scf.select %valid_token, %token0, %c0 : index - %is_neox = index.cmp eq, %mode, %mode_neox : index - %is_normal = index.cmp eq, %mode, %c0 : index - %positions_noalias, %indices_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %input, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer, buffer - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> - %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi64> - %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> - %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> - %index_raw = view.load %indices_view[%token] : view<[%bounded_token_count]xi64> -> i64 - %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 - %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 - %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 - %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 - %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 - %cache_row0 = index.cast %safe_index_i64 : i64 to index - %valid_row = index.cmp ult, %cache_row0, %bounded_cache_row_count : index - %publish1 = scalar.andi %publish0, %valid_index : i1 - %publish = scalar.andi %publish1, %valid_row : i1 - %cache_row = scf.select %valid_row, %cache_row0, %c0 : index - %publish_rotary = scalar.andi %publish, %rotary_channel : i1 - %publish_copy = scalar.andi %publish, %copy_channel : i1 - scf.if %publish_rotary { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %pair_channel = index.mul %channel, %c2 : index - %normal_channel = index.mul %channel, %c4 : index - %input_token_offset = index.mul %token, %input_stride2 : index - %input_head_offset = index.mul %head, %input_stride1 : index - %input_base = index.add %input_token_offset, %input_head_offset : index - %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> - %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> - scf.if %is_neox { - %mscale_vector = vector.splat %mscale : vector<2xf32> - %paired_channel = index.add %pair_channel, %half_n_dims : index - %low_index0 = index.add %input_base, %pair_channel : index - %high_index0 = index.add %input_base, %paired_channel : index - %low_end0 = index.add %low_index0, %c2 : index - %high_end0 = index.add %high_index0, %c2 : index - %low_index, %low_end = index.assume %low_index0, %low_end0 [le(%low_end0, %bounded_input_span)] : index, index - %high_index, %high_end = index.assume %high_index0, %high_end0 [le(%high_end0, %bounded_input_span)] : index, index - %low_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%low_end]xf32> - %high_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%high_end]xf32> - %low_values = vector.load %low_input_view[%low_index] : view<[%low_end]xf32> -> vector<2xf32> - %high_values = vector.load %high_input_view[%high_index] : view<[%high_end]xf32> -> vector<2xf32> - %rotated_low, %rotated_high = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %low_values, %high_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %scaled_low = vector.mulf %rotated_low, %mscale_vector : vector<2xf32> - %scaled_high = vector.mulf %rotated_high, %mscale_vector : vector<2xf32> - template.apply<@ggml.rope_set_rows_f32.store>(%output_format, %publish, %cache_row, %head, %pair_channel, %bounded_cache_row_count, %head_count, %head_size, %scaled_low, %cache_noalias) : (index, i1, index, index, index, index, index, index, vector<2xf32>, buffer) - template.apply<@ggml.rope_set_rows_f32.store>(%output_format, %publish, %cache_row, %head, %paired_channel, %bounded_cache_row_count, %head_count, %head_size, %scaled_high, %cache_noalias) : (index, i1, index, index, index, index, index, index, vector<2xf32>, buffer) - } - scf.if %is_normal { - %mscale_vector = vector.splat %mscale : vector<2xf32> - %normal_index0 = index.add %input_base, %normal_channel : index - %normal_end0 = index.add %normal_index0, %c4 : index - %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index - %normal_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%normal_end]xf32> - %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> - %x0 = vector.extract %values[0] : vector<4xf32> -> f32 - %y0 = vector.extract %values[1] : vector<4xf32> -> f32 - %x1 = vector.extract %values[2] : vector<4xf32> -> f32 - %y1 = vector.extract %values[3] : vector<4xf32> -> f32 - %x_values = vector.from_elements %x0, %x1 : vector<2xf32> - %y_values = vector.from_elements %y0, %y1 : vector<2xf32> - %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %scaled_x = vector.mulf %rotated_x, %mscale_vector : vector<2xf32> - %scaled_y = vector.mulf %rotated_y, %mscale_vector : vector<2xf32> - %scaled_x0 = vector.extract %scaled_x[0] : vector<2xf32> -> f32 - %scaled_x1 = vector.extract %scaled_x[1] : vector<2xf32> -> f32 - %scaled_y0 = vector.extract %scaled_y[0] : vector<2xf32> -> f32 - %scaled_y1 = vector.extract %scaled_y[1] : vector<2xf32> -> f32 - %rotated = vector.from_elements %scaled_x0, %scaled_y0, %scaled_x1, %scaled_y1 : vector<4xf32> - template.apply<@ggml.rope_set_rows_f32.store4>(%output_format, %publish, %cache_row, %head, %normal_channel, %bounded_cache_row_count, %head_count, %head_size, %rotated, %cache_noalias) : (index, i1, index, index, index, index, index, index, vector<4xf32>, buffer) - } - } - scf.if %publish_copy { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %normal_channel = index.mul %channel, %c4 : index - %input_token_offset = index.mul %token, %input_stride2 : index - %input_head_offset = index.mul %head, %input_stride1 : index - %input_base = index.add %input_token_offset, %input_head_offset : index - %normal_index0 = index.add %input_base, %normal_channel : index - %normal_end0 = index.add %normal_index0, %c4 : index - %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index - %normal_input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%normal_end]xf32> - %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> - template.apply<@ggml.rope_set_rows_f32.store4>(%output_format, %publish, %cache_row, %head, %normal_channel, %bounded_cache_row_count, %head_count, %head_size, %values, %cache_noalias) : (index, i1, index, index, index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -kernel.def target(@ggml_rope_set_rows_f32_gfx11_wave32) @ggml_rope_set_rows_f32(%token_count: index, %cache_row_count: index) { - %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index - %head_count = config.get @ggml.rope_set_rows_f32.head_count : index - %head_size = config.get @ggml.rope_set_rows_f32.head_size : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %pair_packet_count = index.div %head_size, %c4 : index - kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index -} launch(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) where [range(%token_count, 1, 2048)] { - template.apply<@ggml.rope_set_rows_f32.body>(%token_count, %cache_row_count, %positions, %indices, %input, %theta, %freq_factors, %cache) : (index, index, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_f32.loom similarity index 100% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/get_rows_f32.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_f32.loom diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/set_rows.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom similarity index 70% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/set_rows.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom index da40e2a1d88c..46bcb16db5a8 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/set_rows.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom @@ -5,7 +5,9 @@ template.decl @ggml.set_rows.launch(%hidden_capacity: index, %token_capacity: in template.decl @ggml.set_rows.load_f32_vector4(%input_format: index, %token: index, %channel: index, %token_count: index, %input_stride: index, %rows: buffer) -> (vector<4xf32>) -template.decl @ggml.set_rows.store_f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) +template.decl @ggml.rows.cache_publish.safe_i64_row(%arg0: index, %arg1: index, %arg2: index, %arg3: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector4(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) amdgpu.target @ggml_set_rows_gfx11_wave64 {subgroup_size = 64} @@ -46,26 +48,6 @@ template.def<@ggml.set_rows.load_f32_vector4> device @ggml_set_rows_load_f32_vec template.return %selected : vector<4xf32> } -template.def<@ggml.set_rows.store_f32_vector4> device @ggml_set_rows_store_f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) { - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c0_offset = index.constant 0 : offset - %is_f16 = index.cmp eq, %output_format, %c16 : index - %is_f32 = index.cmp eq, %output_format, %c32 : index - %publish_f16 = scalar.andi %publish, %is_f16 : i1 - %publish_f32 = scalar.andi %publish, %is_f32 : i1 - scf.if %publish_f16 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf16> - %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> - vector.store %truncated, %cache_view[%cache_row, %channel] : vector<4xf16>, view<[%cache_row_count]x[%hidden_size]xf16> - } - scf.if %publish_f32 { - %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf32> - vector.store %values, %cache_view[%cache_row, %channel] : vector<4xf32>, view<[%cache_row_count]x[%hidden_size]xf32> - } - template.return -} - template.def<@ggml.set_rows.body> device @ggml_set_rows_body(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) { %token_capacity = config.get @ggml.set_rows.token_capacity : index %hidden_capacity = config.get @ggml.set_rows.hidden_capacity : index @@ -81,9 +63,6 @@ template.def<@ggml.set_rows.body> device @ggml_set_rows_body(%token_count: index %c0 = index.constant 0 : index %c4 = index.constant 4 : index %c256 = index.constant 256 : index - %c0_i64 = scalar.constant 0 : i64 - %c1048575_i64 = scalar.constant 1048575 : i64 - %c0_offset = index.constant 0 : offset %packet_tile_base = index.mul %channel_tile, %c256 : index %packet0 = index.add %packet_tile_base, %workitem0 : index %packet_count = index.div %bounded_hidden_size, %c4 : index @@ -94,20 +73,10 @@ template.def<@ggml.set_rows.body> device @ggml_set_rows_body(%token_count: index %packet = scf.select %valid_packet, %packet0, %c0 : index %channel = index.mul %packet, %c4 : index %rows_noalias, %indices_noalias, %cache_noalias = buffer.assume.noalias %rows, %indices, %cache : buffer, buffer, buffer - %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi64> - %index_raw = view.load %indices_view[%token] : view<[%bounded_token_count]xi64> -> i64 - %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 - %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 - %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 - %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 - %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 - %cache_row0 = index.cast %safe_index_i64 : i64 to index - %valid_row = index.cmp ult, %cache_row0, %bounded_cache_row_count : index - %publish1 = scalar.andi %publish0, %valid_index : i1 - %publish = scalar.andi %publish1, %valid_row : i1 - %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %valid_row, %cache_row = template.apply<@ggml.rows.cache_publish.safe_i64_row>(%token, %bounded_token_count, %bounded_cache_row_count, %indices_noalias) : (index, index, index, buffer) -> (i1, index) + %publish = scalar.andi %publish0, %valid_row : i1 %values = template.apply<@ggml.set_rows.load_f32_vector4>(%input_format, %token, %channel, %bounded_token_count, %input_stride, %rows_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>) - template.apply<@ggml.set_rows.store_f32_vector4>(%output_format, %publish, %cache_row, %channel, %bounded_cache_row_count, %bounded_hidden_size, %values, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %channel, %bounded_cache_row_count, %bounded_hidden_size, %values, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) template.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json index 99e80c399c36..be695c991bc2 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json @@ -30,10 +30,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -211,10 +211,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -320,10 +320,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -752,12 +752,15 @@ "token_count": 2 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -921,10 +924,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1081,10 +1084,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1607,12 +1610,15 @@ "token_count": 4 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -1776,10 +1782,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1863,10 +1869,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -2389,12 +2395,15 @@ "token_count": 5 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -2558,10 +2567,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -2694,10 +2703,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -3220,12 +3229,15 @@ "token_count": 3 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -3389,10 +3401,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -3525,10 +3537,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4109,12 +4121,15 @@ "token_count": 16 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -4327,10 +4342,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4435,10 +4450,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4804,12 +4819,15 @@ "token_count": 1 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -5058,10 +5076,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -5471,10 +5489,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json index 3425534725d4..2d739ff260e2 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json @@ -30,10 +30,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -211,10 +211,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -320,10 +320,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -752,12 +752,15 @@ "token_count": 2 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -921,10 +924,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1081,10 +1084,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1607,12 +1610,15 @@ "token_count": 4 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -1776,10 +1782,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -1863,10 +1869,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -2389,12 +2395,15 @@ "token_count": 5 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -2558,10 +2567,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -2694,10 +2703,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -3220,12 +3229,15 @@ "token_count": 3 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -3389,10 +3401,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -3525,10 +3537,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4109,12 +4121,15 @@ "token_count": 16 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -4327,10 +4342,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4435,10 +4450,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -4804,12 +4819,15 @@ "token_count": 1 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -5058,10 +5076,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -5491,10 +5509,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json index 521002c7d7ee..e9e3eb8b7c04 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json @@ -30,10 +30,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -184,10 +184,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -266,10 +266,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", @@ -674,12 +674,15 @@ "token_count": 4 }, "kernel": "loom_libs:ggml_set_rows", - "library_sources": [], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], "primary_sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom" ], "sources": [ - "ops/set_rows.loom" + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" ], "symbol": "ggml_set_rows", "workload_parameters": [ @@ -843,10 +846,10 @@ "motifs/f32_f16.loom" ], "primary_sources": [ - "ops/get_rows_f32.loom" + "ops/rows/get_rows_f32.loom" ], "sources": [ - "ops/get_rows_f32.loom", + "ops/rows/get_rows_f32.loom", "motifs/dequant.loom", "motifs/publish_f32.loom", "motifs/quantize_q8_1_x4.loom", From ffab34aaa92177a1abdc186387cd3f2d4ae03b0e Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 13:35:13 -0700 Subject: [PATCH 5/8] Add parallel HRX test suite runner Introduce a reusable test suite utility for process-isolated parallel test execution, output controls, listing, filtering, and device-aware skips. Convert the HRX ops, buffer, backend, and Loom JIT tests to register named cases with the shared suite runner. Assisted-by: Codex --- tests/hrx-backend-test.cpp | 175 +++++---- tests/test-hrx-buffer.cpp | 53 ++- tests/test-hrx-loom-jit.cpp | 38 +- tests/test-hrx-ops.cpp | 716 +++++++++++++++++++++++------------- tests/testing_suite.h | 548 +++++++++++++++++++++++++++ 5 files changed, 1172 insertions(+), 358 deletions(-) create mode 100644 tests/testing_suite.h diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index cd4c27766e79..cf3761120ac8 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -23,6 +23,7 @@ #include "runtime/graph-program-cache.h" #include "runtime/graph-replay.h" #include "runtime/loom-kernel-jit.h" +#include "testing_suite.h" #include #include @@ -10647,85 +10648,97 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_free(ctx); } -int main() { - run_status_checks(); - run_command_plan_metadata_checks(); - run_dispatch_registry_checks(); - run_graph_import_checks(); - run_graph_import_mixed_backend_boundary_checks(); - run_graph_view_external_use_checks(); - run_scale_f32_dispatch_checks(); - run_cont_f32_dispatch_checks(); - run_binary_f32_broadcast_dispatch_checks(); - run_graph_snapshot_diagnostics_checks(); - run_unmatched_graph_diagnostics_checks(); - run_completion_counter_plan_checks(); - run_graph_index_checks(); - run_rope_set_rows_dispatch_checks(); - run_graph_traversal_checks(); - run_qwen_token_embedding_dispatch_checks(); - run_gather_add_dispatch_checks(); - run_qwen_flash_attention_dispatch_checks(); - run_tiled_pair_matmul_postops_dispatch_checks(); - run_qwen_attention_postprocess_dispatch_checks(); - run_qwen_matmul_dispatch_checks(); - run_packed_f16_producer_consumer_checks(); - run_tiled_matmul_alternate_publish_checks(); - run_k16_major_preparation_reuse_checks(); - run_rmsnorm_binary_k16_output_checks(); - run_swiglu_q8_output_checks(); - run_rmsnorm_gate_packed_output_checks(); - run_rmsnorm_gate_q8_output_checks(); - run_lowtoken_residual_dispatch_checks(); - run_quantized_conv4_dispatch_checks(); - run_gdn_rmsnorm_gate_dispatch_checks(); - run_gdn_native_projection_pair_dispatch_checks(); - run_gdn_selected_snapshot_dispatch_checks(); - run_gdn_selected_rms_q8_dispatch_checks(); - run_llama_attention_matmul_dispatch_checks(); - run_quantized_value_projection_dispatch_checks(); - schedule_qwen_terminal_q6k_q8_command(1); - schedule_qwen_terminal_q6k_q8_command(18); - schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); - schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); - run_qwen_router_top8_dispatch_checks(); - run_common_mul_mat_id_swiglu_dispatch_checks(); - run_qwen_routed_gate_up_dispatch_checks(); - run_common_mul_mat_id_postops_dispatch_checks(); - run_alias_value_import_checks(); - run_multi_dispatch_checks(); - run_layout_alias_scheduler_elision_checks(); - run_zero_output_scheduler_elision_checks(); - run_transient_import_checks(); - run_graph_view_preserves_shared_output_storage_checks(); - run_chained_dispatch_requires_transients(); - run_graph_replay_host_staging_is_not_ineligible(); - run_multiple_transient_plan_checks(); - run_disjoint_transient_plan_packing_checks(); - run_command_program_kernel_dump_checks(); - run_command_shape_hash_checks(); - run_graph_program_cache_uid_mismatch_checks(); - run_graph_match_hash_collision_checks(); - run_graph_match_bijection_checks(); - run_validated_graph_uid_match_checks(); - run_graph_executor_contract_checks(); - REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); - - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); - return 0; - } - - run_zero_output_device_support_checks(); - run_scale_f32_device_support_checks(); - run_qwen_expert_table_partition_prefill_512_execution(); - run_add_f32(); - run_scale_f32(); - run_scale_f32_inplace(); - run_two_independent_add_f32(); - run_chained_add_f32(); - run_same_uid_distinct_graph_reuses_graph_program(); - run_same_uid_distinct_graph_reuses_graph_program(true); - run_unsupported_op_fails(); - return 0; +static void register_hrx_backend_host_cases(test_runner::Suite & suite) { + suite.host_case("status", [] { run_status_checks(); }); + suite.host_case("command_plan_metadata", [] { run_command_plan_metadata_checks(); }); + suite.host_case("dispatch_registry", [] { run_dispatch_registry_checks(); }); + suite.host_case("graph_import", [] { run_graph_import_checks(); }); + suite.host_case("graph_import_mixed_backend_boundary", [] { run_graph_import_mixed_backend_boundary_checks(); }); + suite.host_case("graph_view_external_use", [] { run_graph_view_external_use_checks(); }); + suite.host_case("scale_f32_dispatch", [] { run_scale_f32_dispatch_checks(); }); + suite.host_case("cont_f32_dispatch", [] { run_cont_f32_dispatch_checks(); }); + suite.host_case("binary_f32_broadcast_dispatch", [] { run_binary_f32_broadcast_dispatch_checks(); }); + suite.host_case("graph_snapshot_diagnostics", [] { run_graph_snapshot_diagnostics_checks(); }); + suite.host_case("unmatched_graph_diagnostics", [] { run_unmatched_graph_diagnostics_checks(); }); + suite.host_case("completion_counter_plan", [] { run_completion_counter_plan_checks(); }); + suite.host_case("graph_index", [] { run_graph_index_checks(); }); + suite.host_case("rope_set_rows_dispatch", [] { run_rope_set_rows_dispatch_checks(); }); + suite.host_case("graph_traversal", [] { run_graph_traversal_checks(); }); + suite.host_case("qwen_token_embedding_dispatch", [] { run_qwen_token_embedding_dispatch_checks(); }); + suite.host_case("gather_add_dispatch", [] { run_gather_add_dispatch_checks(); }); + suite.host_case("qwen_flash_attention_dispatch", [] { run_qwen_flash_attention_dispatch_checks(); }); + suite.host_case("tiled_pair_matmul_postops_dispatch", [] { run_tiled_pair_matmul_postops_dispatch_checks(); }); + suite.host_case("qwen_attention_postprocess_dispatch", [] { run_qwen_attention_postprocess_dispatch_checks(); }); + suite.host_case("qwen_matmul_dispatch", [] { run_qwen_matmul_dispatch_checks(); }); + suite.host_case("packed_f16_producer_consumer", [] { run_packed_f16_producer_consumer_checks(); }); + suite.host_case("tiled_matmul_alternate_publish", [] { run_tiled_matmul_alternate_publish_checks(); }); + suite.host_case("k16_major_preparation_reuse", [] { run_k16_major_preparation_reuse_checks(); }); + suite.host_case("rmsnorm_binary_k16_output", [] { run_rmsnorm_binary_k16_output_checks(); }); + suite.host_case("swiglu_q8_output", [] { run_swiglu_q8_output_checks(); }); + suite.host_case("rmsnorm_gate_packed_output", [] { run_rmsnorm_gate_packed_output_checks(); }); + suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); + suite.host_case("lowtoken_residual_dispatch", [] { run_lowtoken_residual_dispatch_checks(); }); + suite.host_case("quantized_conv4_dispatch", [] { run_quantized_conv4_dispatch_checks(); }); + suite.host_case("gdn_rmsnorm_gate_dispatch", [] { run_gdn_rmsnorm_gate_dispatch_checks(); }); + suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); + suite.host_case("gdn_selected_snapshot_dispatch", [] { run_gdn_selected_snapshot_dispatch_checks(); }); + suite.host_case("gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); }); + suite.host_case("llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); }); + suite.host_case("quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); }); + suite.host_case("qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); }); + suite.host_case("qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); }); + suite.host_case("get_rows_q8_1_alternate_command.q4_k", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); }); + suite.host_case("get_rows_q8_1_alternate_command.bf16", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); }); + suite.host_case("qwen_router_top8_dispatch", [] { run_qwen_router_top8_dispatch_checks(); }); + suite.host_case("common_mul_mat_id_swiglu_dispatch", [] { run_common_mul_mat_id_swiglu_dispatch_checks(); }); + suite.host_case("qwen_routed_gate_up_dispatch", [] { run_qwen_routed_gate_up_dispatch_checks(); }); + suite.host_case("common_mul_mat_id_postops_dispatch", [] { run_common_mul_mat_id_postops_dispatch_checks(); }); + suite.host_case("alias_value_import", [] { run_alias_value_import_checks(); }); + suite.host_case("multi_dispatch", [] { run_multi_dispatch_checks(); }); + suite.host_case("layout_alias_scheduler_elision", [] { run_layout_alias_scheduler_elision_checks(); }); + suite.host_case("zero_output_scheduler_elision", [] { run_zero_output_scheduler_elision_checks(); }); + suite.host_case("transient_import", [] { run_transient_import_checks(); }); + suite.host_case("graph_view_preserves_shared_output_storage", [] { run_graph_view_preserves_shared_output_storage_checks(); }); + suite.host_case("chained_dispatch_requires_transients", [] { run_chained_dispatch_requires_transients(); }); + suite.host_case("graph_replay_host_staging_is_not_ineligible", [] { run_graph_replay_host_staging_is_not_ineligible(); }); + 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_shape_hash", [] { run_command_shape_hash_checks(); }); + suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); + suite.host_case("graph_match_hash_collision", [] { run_graph_match_hash_collision_checks(); }); + suite.host_case("graph_match_bijection", [] { run_graph_match_bijection_checks(); }); + suite.host_case("validated_graph_uid_match", [] { run_validated_graph_uid_match_checks(); }); + suite.host_case("graph_executor_contract", [] { run_graph_executor_contract_checks(); }); + suite.host_case("loom_async_jit_environment", [] { + REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); + }); +} + +static void register_hrx_backend_device_cases(test_runner::Suite & suite) { + suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); + suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); + suite.device_case("qwen_expert_table_partition_prefill_512_execution", [] { run_qwen_expert_table_partition_prefill_512_execution(); }); + suite.device_case("add_f32", [] { run_add_f32(); }); + suite.device_case("scale_f32", [] { run_scale_f32(); }); + suite.device_case("scale_f32_inplace", [] { run_scale_f32_inplace(); }); + suite.device_case("two_independent_add_f32", [] { run_two_independent_add_f32(); }); + suite.device_case("chained_add_f32", [] { run_chained_add_f32(); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.default", [] { run_same_uid_distinct_graph_reuses_graph_program(); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.distinct_uid", [] { run_same_uid_distinct_graph_reuses_graph_program(true); }); + suite.device_case("unsupported_op_fails", [] { run_unsupported_op_fails(); }); +} + +static void register_hrx_backend_cases(test_runner::Suite & suite) { + register_hrx_backend_host_cases(suite); + register_hrx_backend_device_cases(suite); +} + +int main(int argc, char ** argv) { + test_runner::Suite suite(test_runner::Config::with_prefix( + "HRX backend test", "hrx-backend", "GGML_HRX_BACKEND_TEST")); + register_hrx_backend_cases(suite); + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + return suite.run(argc, argv, has_device); } diff --git a/tests/test-hrx-buffer.cpp b/tests/test-hrx-buffer.cpp index 75a27ea4e5e3..e281a3199f09 100644 --- a/tests/test-hrx-buffer.cpp +++ b/tests/test-hrx-buffer.cpp @@ -6,6 +6,7 @@ #include "ggml.h" #include "hrx-interop-utils.h" #include "runtime/host-memory.h" +#include "testing_suite.h" #include #include @@ -352,22 +353,48 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(weight_stats.resident_bytes == 0); } -int main() { - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); - return 0; - } - +template +static void run_with_hrx_backend(F && f) { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); ggml_backend_hrx_context * context = backend_context(backend); + f(backend, context); + ggml_backend_free(backend); +} - run_backend_buffer_checks(backend); - run_host_buffer_checks(backend); - run_host_transfer_checks(context); - run_host_staging_checks(context); - run_host_weight_cache_checks(context); +static void register_hrx_buffer_cases(test_runner::Suite & suite) { + suite.device_case("backend_buffer", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_backend_buffer_checks(backend); + }); + }); + suite.device_case("host_buffer", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_host_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); + }); + }); + suite.device_case("host_staging", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_host_staging_checks(context); + }); + }); + suite.device_case("host_weight_cache", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_host_weight_cache_checks(context); + }); + }); +} - ggml_backend_free(backend); - return 0; +int main(int argc, char ** argv) { + test_runner::Suite suite(test_runner::Config::with_prefix( + "HRX buffer test", "hrx-buffer", "GGML_HRX_BUFFER_TEST", 1)); + register_hrx_buffer_cases(suite); + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + return suite.run(argc, argv, has_device); } diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index 880ce1f56b86..ff8c376c3bc4 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -3,6 +3,7 @@ #include "kernel-corpus/kernel-corpus.h" #include "runtime/kernel-executable-cache.h" #include "runtime/loom-kernel-jit.h" +#include "testing_suite.h" #include #include @@ -446,7 +447,7 @@ static void run_targeted_export_materialize_case(HrxTestDevice & device) { } // namespace -int main() { +static void run_loom_jit_compile_checks() { static constexpr const char * kTarget = "gfx1100"; const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); @@ -1006,14 +1007,33 @@ int main() { compiled.reset(); } +} + +static void run_loom_jit_materialize_checks() { + const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); + HrxTestDevice device; - if (device.open()) { - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); - run_targeted_export_materialize_case(device); - } else { - std::printf("skipping KernelExecutableCache materialization checks: no HRX device available\n"); - } + REQUIRE(device.open()); + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); + run_targeted_export_materialize_case(device); +} + +static bool has_hrx_test_device() { + HrxTestDevice device; + return device.open(); +} + +static void register_hrx_loom_jit_cases(test_runner::Suite & suite) { + suite.host_case("compile", [] { run_loom_jit_compile_checks(); }); + suite.device_case("materialize", [] { run_loom_jit_materialize_checks(); }); +} + +int main(int argc, char ** argv) { + test_runner::Suite suite(test_runner::Config::with_prefix( + "HRX Loom JIT test", "hrx-loom-jit", "GGML_HRX_LOOM_JIT_TEST", 2)); + register_hrx_loom_jit_cases(suite); - return 0; + const bool has_device = has_hrx_test_device(); + return suite.run(argc, argv, has_device); } diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index f6a2219381b7..c47f46a2ecd2 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -7,6 +7,7 @@ #include "graph/graph.h" #include "kernel-corpus/kernel-corpus.h" #include "runtime/graph-executor.h" +#include "testing_suite.h" #include #include @@ -4605,285 +4606,490 @@ static void run_vector_postops_cpu_reference_case() { false, kVectorInputSize); } -int main() { - const bool run_vector_cap_experiment = std::getenv("GGML_HRX_RUN_VECTOR_CAP_EXPERIMENT") != nullptr; - const bool run_vector_publish_test = std::getenv("GGML_HRX_RUN_VECTOR_PUBLISH_TEST") != nullptr; - const bool run_vector_postops_test = std::getenv("GGML_HRX_RUN_VECTOR_POSTOPS_TEST") != nullptr; - if (!run_vector_cap_experiment && !run_vector_publish_test && !run_vector_postops_test) { - run_rmsnorm_support_checks(); - run_alternate_value_alias_lookup_checks(); +static const char * glu_op_name(ggml_glu_op op) { + switch (op) { + case GGML_GLU_OP_REGLU: + return "reglu"; + case GGML_GLU_OP_GEGLU: + return "geglu"; + case GGML_GLU_OP_SWIGLU: + return "swiglu"; + case GGML_GLU_OP_SWIGLU_OAI: + return "swiglu_oai"; + case GGML_GLU_OP_GEGLU_ERF: + return "geglu_erf"; + case GGML_GLU_OP_GEGLU_QUICK: + return "geglu_quick"; + case GGML_GLU_OP_COUNT: + break; } + return "unknown_glu"; +} - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); - return 0; - } +static std::string type_name(ggml_type type) { + return ggml_type_name(type); +} - if (run_vector_cap_experiment) { - run_vector_output_capacity_experiment(); - return 0; +using test_runner::Suite; + +static void register_basic_ops_cases(Suite & suite) { + suite.host_case("support.rmsnorm", [] { run_rmsnorm_support_checks(); }); + suite.host_case("support.alternate_value_alias_lookup", [] { run_alternate_value_alias_lookup_checks(); }); + suite.device_case("add_f32.cpu_reference", [] { run_add_f32_cpu_reference_case(); }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".plain", + [type] { run_ssm_conv_prefill_recurrent_case(type, false, false); }); + suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered", + [type] { run_ssm_conv_prefill_recurrent_case(type, true, false); }); + suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered_clear_state", + [type] { run_ssm_conv_prefill_recurrent_case(type, true, true); }); } - if (run_vector_publish_test) { - run_vector_q8_publish_cpu_reference_case(); - return 0; + suite.device_case("view_mul_f32.cpu_reference", [] { run_view_mul_f32_cpu_reference_case(); }); + for (const int64_t token_count : { 2, 23 }) { + suite.device_case("strided_cont.tokens" + std::to_string(token_count), + [token_count] { run_strided_cont_cpu_reference_case(token_count); }); } - if (run_vector_postops_test) { - run_vector_postops_cpu_reference_case(); - return 0; + for (const ggml_glu_op op : + { GGML_GLU_OP_REGLU, GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { + suite.device_case(std::string("glu_split_f32.") + glu_op_name(op), + [op] { run_glu_split_f32_cpu_reference_case(op); }); + suite.device_case(std::string("glu_packed_f32.") + glu_op_name(op) + ".normal", + [op] { run_glu_packed_f32_cpu_reference_case(op, false); }); } - - run_add_f32_cpu_reference_case(); - for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - run_ssm_conv_prefill_recurrent_case(type, false, false); - run_ssm_conv_prefill_recurrent_case(type, true, false); - run_ssm_conv_prefill_recurrent_case(type, true, true); + suite.device_case("glu_packed_f32.swiglu.gated", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, true); }); + suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens2", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 2); }); + suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens14", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); }); + suite.device_case("gather_add_f32.cpu_reference", [] { run_gather_add_f32_cpu_reference_case(); }); + for (const ggml_type type : + { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { + suite.device_case("get_rows_f32." + type_name(type), [type] { run_get_rows_f32_cpu_reference_case(type); }); } - run_view_mul_f32_cpu_reference_case(); - run_strided_cont_cpu_reference_case(2); - run_strided_cont_cpu_reference_case(23); - run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_REGLU); - run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU); - run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU); - run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_ERF); - run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_QUICK); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_REGLU, false); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU, false); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_ERF, false); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_QUICK, false); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, true); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 2); - run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); - run_gather_add_f32_cpu_reference_case(); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q4_K); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q6_K); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q8_0); - run_get_rows_q8_1_zero_weight_case(); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_F16); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_BF16); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32); - for (int64_t tokens : { 1, 3, 5 }) { - for (int64_t outputs : { 1, 4, 47, 48, 63 }) { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_f32_f32_wmma", - tokens, outputs); + suite.device_case("get_rows_f32.q1_0.rows2048", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); }); + suite.device_case("get_rows_f32.q5_1.rows640", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); + suite.device_case("get_rows_q8_1.zero_weight", [] { run_get_rows_q8_1_zero_weight_case(); }); +} + +static void register_dense_matmul_cases(Suite & suite) { + for (const int64_t tokens : { 1, 3, 5 }) { + for (const int64_t outputs : { 1, 4, 47, 48, 63 }) { + suite.device_case( + "dense_matmul.q4_k.tokens" + std::to_string(tokens) + ".outputs" + std::to_string(outputs), + [tokens, outputs] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, outputs); + }); } } - for (ggml_type type : { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, - GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + for (const ggml_type type : + { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { + suite.device_case("dense_matmul." + type_name(type) + ".tokens3.outputs47", [type] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + }); } - run_dense_matmul_unary_cpu_reference_case( - GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); - run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - for (ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { - run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32", - 1, 1024, 2048, false, false, cache_type); + suite.device_case("dense_matmul_unary.f32.tokens3.outputs47", [] { + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + }); + suite.device_case("get_rows_f32.f32.rows262144.tokens4.outputs1", + [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); }); + suite.device_case("dense_matmul.q4_k.wmma.tokens2.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + }); + suite.device_case("dense_matmul.q4_k.vector.tokens1.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + for (const ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { + suite.device_case( + "attention_v_matmul." + type_name(type) + ".cache_" + type_name(cache_type), [type, cache_type] { + run_dense_matmul_cpu_reference_case(type, + "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32", 1, + 1024, 2048, false, false, cache_type); + }); } } - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128, 544); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 1024, 640); - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 1024, 640); - run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_wmma"); - run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU); - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 2048, GGML_GLU_OP_GEGLU, - 640); - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_REGLU); - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU_ERF); - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU_QUICK); - run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128); - run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128, - GGML_GLU_OP_SWIGLU, true); - run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); - run_vector_q8_publish_cpu_reference_case(); - for (ggml_glu_op op : { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, - GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128, op); - 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, 4096, 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", 256, 128, 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", 256, 4096, 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); + 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_Q6_K ? "loom_libs:ggml_mul_mat_vector_q6_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; + suite.device_case("dense_matmul." + type_name(type) + ".wmma.tokens2.outputs" + std::to_string(outputs), + [type, outputs, input] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, + outputs, input); + }); + suite.device_case("dense_matmul." + type_name(type) + ".vector.tokens1.outputs" + std::to_string(outputs), + [type, vector_kernel, outputs, input] { + run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); + }); + } + suite.device_case("dense_matmul.q5_0.wmma.tokens2.outputs256.input640", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256, 640); + }); + suite.device_case("dense_matmul.q5_0.wmma.tokens2.outputs128.input544", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128, 544); + }); + suite.device_case("dense_matmul.q5_0.vector.tokens1.outputs256.input640", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); + }); + suite.device_case("dense_matmul.q8_1.zero_weight.wmma", + [] { run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_wmma"); }); + suite.device_case("dense_matmul.q8_1.zero_weight.vector", [] { + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); + }); + suite.device_case("dense_matmul.f32.wmma.tokens2.outputs256", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); + }); + suite.device_case("dense_matmul.f32.vector.tokens1.outputs256", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); + }); + suite.device_case("dense_matmul_unary.f32.tokens2.outputs128", [] { + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + }); + suite.device_case("dense_matmul_swiglu.q4_k.f16.swiglu.tokens2.outputs128", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + }); + suite.device_case("dense_matmul_swiglu.q4_k.f16.geglu.tokens2.outputs128", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, + GGML_GLU_OP_GEGLU); + }); + suite.device_case("dense_matmul_swiglu.iq4_nl.iq4_nl.geglu.tokens2.outputs2048.input640", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, + 2048, GGML_GLU_OP_GEGLU, 640); + }); + for (const ggml_glu_op op : { GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { + suite.device_case(std::string("dense_matmul_swiglu.q4_k.f16.") + glu_op_name(op) + ".tokens2.outputs128", [op] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, op); + }); + } + suite.device_case("dense_matmul_packed_glu.q4_k.tokens2.outputs128", [] { + run_dense_matmul_packed_glu_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128); + }); + suite.device_case("dense_matmul_packed_glu.q4_k.swiglu.tokens2.outputs128.gated", [] { + run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + 2, 128, GGML_GLU_OP_SWIGLU, true); + }); + suite.device_case("dense_matmul_binary.q4_k.f16.tokens2.outputs128", [] { + run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + }); + suite.device_case("vector_q8_publish.cpu_reference", [] { run_vector_q8_publish_cpu_reference_case(); }); + for (const ggml_glu_op op : + { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { + for (const int64_t outputs : { 128, 4096 }) { + suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.lowtoken.") + glu_op_name(op) + + ".tokens4.outputs" + std::to_string(outputs), + [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); + }); + } + for (const int64_t tokens : { 256, 512 }) { + if (tokens == 512) { + suite.device_case( + std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + ".tokens512.outputs4096", + [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); + }); + } else { + for (const int64_t outputs : { 128, 4096 }) { + suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + + ".tokens256.outputs" + std::to_string(outputs), + [op, outputs] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + outputs, op); + }); + } + } + } } for (const int64_t tokens : { 1, 2, 3, 5 }) { - 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); + suite.device_case( + "dense_matmul_swiglu.q4_k.q4_k.lowtoken.tokens" + std::to_string(tokens) + ".outputs4160", [tokens] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 4160); + }); } for (const int64_t tokens : { 1, 2, 5 }) { - 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); + suite.device_case( + "dense_matmul_swiglu.q4_k.q4_k.swiglu.tokens" + std::to_string(tokens) + ".outputs16384.input4096", + [tokens] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); + }); } - 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); - 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); - 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); - 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); - 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); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", 33, 256, - true, false, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, - false, true, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", 33, - 256, true, true, false); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, - false, true, true); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", 33, - 256, true, true, true); - for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; - run_dense_matmul_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, input_size); + 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); + }); + 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); + }); + 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); + }); + 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); + }); + 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); + }); + suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32", 33, 256, true, false, false); + }); + suite.device_case("dense_matmul_postops.f16.residual.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, true, false, input_size + 256); + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, false); + }); + suite.device_case("dense_matmul_postops.f16.bias_residual.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, true, false); + }); + suite.device_case("dense_matmul_postops.f16.residual.alias.tokens33.outputs256", [] { run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, false, true, true, input_size + 256); + GGML_TYPE_F16, "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32", 33, 256, false, true, true); + }); + suite.device_case("dense_matmul_postops.f16.bias_residual.alias.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, true, true); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; + suite.device_case( + "dense_matmul." + type_name(type) + ".wmma.tokens5.outputs4096.input" + std::to_string(input_size), + [type, input_size] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, input_size); + }); + suite.device_case("dense_matmul_postops." + type_name(type) + ".add.tokens5.outputs4160", [type, input_size] { + run_dense_matmul_postops_cpu_reference_case(type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, + true, false, input_size + 256); + }); + suite.device_case( + "dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", [type, input_size] { + run_dense_matmul_postops_cpu_reference_case(type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, + false, true, true, input_size + 256); + }); } - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 8192); + suite.device_case("dense_matmul.q4_k.wmma.tokens5.outputs4096.input8192", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 8192); + }); for (const int64_t input_size : { 6144, 6400, 6656 }) { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 5120, input_size); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 5120, false, true, true, input_size); + suite.device_case("dense_matmul.q4_k.wmma.tokens5.outputs5120.input" + std::to_string(input_size), + [input_size] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", + 5, 5120, input_size); + }); + suite.device_case( + "dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", + 5, 5120, false, true, true, input_size); + }); } - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 4096, 4096); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4160, false, true, false, 6400); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4224, false, true, true, 6656); - run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 16384); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, true, false, 16640); - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, false, true, true, 16640); + suite.device_case("dense_matmul.q4_k.wmma.tokens4.outputs4096.input4096", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 4096, 4096); + }); + suite.device_case("dense_matmul_postops.q4_k.add.tokens4.outputs4160.input6400", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4160, + false, true, false, 6400); + }); + suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens4.outputs4224.input6656", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4224, + false, true, true, 6656); + }); + suite.device_case("endpoint_rmsnorm_q6k_q8.cpu_reference", + [] { run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); }); + suite.device_case("dense_matmul.q6_k.wmma.tokens5.outputs4096.input16384", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 16384); + }); + suite.device_case("dense_matmul_postops.q6_k.add.tokens5.outputs4160.input16640", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, + false, true, false, 16640); + }); + suite.device_case("dense_matmul_postops.q6_k.add.alias.tokens5.outputs4224.input16640", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, + false, true, true, 16640); + }); for (const int64_t tokens : { 1, 3, 5 }) { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_f32_f32_wmma", - tokens, 256, 5120, true); - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_f32_f32_wmma", - tokens, 256, 5120, true, true); + suite.device_case( + "dense_matmul.q4_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, 256, 5120, true); + }); + suite.device_case( + "dense_matmul.q6_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q6_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, 256, 5120, true, true); + }); } - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 256, 5120, true); - run_rmsnorm_cpu_reference_case(); - run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); - run_strided_rmsnorm_cpu_reference_case(2); - run_strided_rmsnorm_cpu_reference_case(23); - run_rmsnorm_binary_add_cpu_reference_case(); - run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); - run_rmsnorm_gate_cpu_reference_case(256, 3, 3, GGML_UNARY_OP_GELU); - run_rmsnorm_gate_cpu_reference_case(512, 3, 5, GGML_UNARY_OP_RELU); - run_rmsnorm_gate_cpu_reference_case(1024, 1, 4, GGML_UNARY_OP_GELU_ERF); - run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); - run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); - run_gemma_scaled_rmsnorm_cpu_reference_case(); - run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); - run_scheduled_hrx_scale_rmsnorm_case(); - run_scheduled_hrx_rmsnorm_scale_case(); - run_scheduled_hrx_rmsnorm_view_scale_case(); - run_external_view_input_rmsnorm_case(); - run_split_local_view_alias_import_case(); - run_rope_set_rows_cpu_reference_case(); - run_attention_postprocess_cpu_reference_case(); - run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); - run_asymmetric_flash_attention_cpu_reference_case(); - run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); - run_decode_attention_qkv_scheduling_case(); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); - run_rmsnorm_mul_case(256, 1); - run_rmsnorm_mul_case(256, 4); - run_rmsnorm_mul_case(2048, 1); - run_router_projection_case(4); - run_router_top8_case(4); - run_qwen_flash_attention_case(); + 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, + true); + }); +} + +static void register_rmsnorm_and_scheduling_cases(Suite & suite) { + suite.device_case("rmsnorm.default", [] { run_rmsnorm_cpu_reference_case(); }); + suite.device_case("rmsnorm.hidden3840.tokens18", [] { run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); }); + for (const int64_t token_count : { 2, 23 }) { + suite.device_case("strided_rmsnorm.tokens" + std::to_string(token_count), + [token_count] { run_strided_rmsnorm_cpu_reference_case(token_count); }); + } + suite.device_case("rmsnorm_binary_add.cpu_reference", [] { run_rmsnorm_binary_add_cpu_reference_case(); }); + suite.device_case("rmsnorm_gate.silu.hidden128.tokens1.stride7", + [] { run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); }); + suite.device_case("rmsnorm_gate.gelu.hidden256.tokens3.stride3", + [] { run_rmsnorm_gate_cpu_reference_case(256, 3, 3, GGML_UNARY_OP_GELU); }); + suite.device_case("rmsnorm_gate.relu.hidden512.tokens3.stride5", + [] { run_rmsnorm_gate_cpu_reference_case(512, 3, 5, GGML_UNARY_OP_RELU); }); + suite.device_case("rmsnorm_gate.gelu_erf.hidden1024.tokens1.stride4", + [] { run_rmsnorm_gate_cpu_reference_case(1024, 1, 4, GGML_UNARY_OP_GELU_ERF); }); + suite.device_case("rmsnorm_gate.silu.hidden128.tokens48.stride512", + [] { run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); }); + suite.device_case("rmsnorm_mul.hidden3840.tokens18", [] { run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); }); + suite.device_case("gemma_scaled_rmsnorm.cpu_reference", [] { run_gemma_scaled_rmsnorm_cpu_reference_case(); }); + suite.device_case("gemma_scaled_rmsnorm_mul.cpu_reference", + [] { run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); }); + suite.device_case("scheduled.scale_rmsnorm", [] { run_scheduled_hrx_scale_rmsnorm_case(); }); + suite.device_case("scheduled.rmsnorm_scale", [] { run_scheduled_hrx_rmsnorm_scale_case(); }); + suite.device_case("scheduled.rmsnorm_view_scale", [] { run_scheduled_hrx_rmsnorm_view_scale_case(); }); + suite.device_case("external_view_input_rmsnorm", [] { run_external_view_input_rmsnorm_case(); }); + suite.device_case("split_local_view_alias_import", [] { run_split_local_view_alias_import_case(); }); + suite.device_case("rope_set_rows.cpu_reference", [] { run_rope_set_rows_cpu_reference_case(); }); + suite.device_case("attention_postprocess.cpu_reference", [] { run_attention_postprocess_cpu_reference_case(); }); + suite.device_case("routed_moe.q4_k.next_rmsnorm", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); }); + suite.device_case("flash_attention.asymmetric", [] { run_asymmetric_flash_attention_cpu_reference_case(); }); + suite.device_case("routed_moe.q6_k", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); }); + suite.device_case("decode_attention_qkv.scheduling", [] { run_decode_attention_qkv_scheduling_case(); }); + suite.device_case("decode_routed_moe.q4_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); }); + suite.device_case("decode_routed_moe.q6_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); }); + suite.device_case("decode_routed_moe.q6_k.alias_gate_input", + [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); }); + suite.device_case("rmsnorm_mul.hidden256.tokens1", [] { run_rmsnorm_mul_case(256, 1); }); + 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("router_top8.tokens4", [] { run_router_top8_case(4); }); + suite.device_case("qwen_flash_attention", [] { run_qwen_flash_attention_case(); }); +} + +static void register_flash_attention_cases(Suite & suite) { for (int64_t head_size = 64; head_size <= 512; head_size += 64) { - run_common_flash_attention_cpu_reference_case(head_size, 2, 65, - 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - run_common_flash_attention_cpu_reference_case(head_size, 16, 64, - 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); - run_common_flash_attention_cpu_reference_case(head_size, 17, 65, - 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); + suite.device_case("common_flash_attention.decode.head" + std::to_string(head_size) + ".query2.kv65", + [head_size] { + run_common_flash_attention_cpu_reference_case( + head_size, 2, 65, 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + }); + suite.device_case( + "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query16.kv64", [head_size] { + run_common_flash_attention_cpu_reference_case(head_size, 16, 64, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); + }); + suite.device_case( + "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query17.kv65", [head_size] { + run_common_flash_attention_cpu_reference_case(head_size, 17, 65, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); + }); + } + suite.device_case("common_flash_attention.decode.head96.query2.kv4.capacity320", [] { + run_common_flash_attention_cpu_reference_case( + 96, 2, 4, 1.0f / std::sqrt(96.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 320); + }); + suite.device_case("common_flash_attention.decode.head320.query2.kv64.capacity64", [] { + run_common_flash_attention_cpu_reference_case( + 320, 2, 64, 1.0f / std::sqrt(320.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + 64); + }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512", + [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512); }); + suite.device_case("qwen_decode_split_flash_attention.query4.kv513", + [] { run_qwen_decode_split_flash_attention_scheduling_case(4, 513); }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked", + [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true); }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked.head96.value64", [] { + run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true, 96, 64, 1.0f / std::sqrt(96.0f)); + }); + suite.device_case("qwen_decode_split_flash_attention.query4.kv513.masked.head96.value64", [] { + run_qwen_decode_split_flash_attention_scheduling_case(4, 513, true, 96, 64, 1.0f / std::sqrt(96.0f)); + }); + suite.device_case("qwen_decode_attention_output_next_q8.no_selectors", + [] { run_qwen_decode_attention_output_next_q8_scheduling_case(false); }); + suite.device_case("qwen_decode_attention_output_next_q8.selectors", + [] { run_qwen_decode_attention_output_next_q8_scheduling_case(true); }); + suite.device_case("qwen_full_cache_prefill_flash_attention.tokens512", + [] { run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); }); +} + +static void register_hrx_ops_cases(Suite & suite) { + register_basic_ops_cases(suite); + register_dense_matmul_cases(suite); + register_rmsnorm_and_scheduling_cases(suite); + register_flash_attention_cases(suite); +} + +int main(int argc, char ** argv) { + Suite suite(test_runner::Config::with_prefix("HRX test", "hrx", "GGML_HRX_TEST")); + register_hrx_ops_cases(suite); + + test_runner::Options options; + if (!suite.parse_options(argc, argv, options)) { + suite.print_usage(argc > 0 ? argv[0] : "test-hrx-ops"); + return 2; + } + if (options.help) { + return suite.run(options, false); } - run_common_flash_attention_cpu_reference_case(96, 2, 4, 1.0f / std::sqrt(96.0f), - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 320); - run_common_flash_attention_cpu_reference_case(320, 2, 64, 1.0f / std::sqrt(320.0f), - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 64); - run_qwen_decode_split_flash_attention_scheduling_case(1, 512); - run_qwen_decode_split_flash_attention_scheduling_case(4, 513); - run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true); - run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true, 96, 64, 1.0f / std::sqrt(96.0f)); - run_qwen_decode_split_flash_attention_scheduling_case(4, 513, true, 96, 64, 1.0f / std::sqrt(96.0f)); - run_qwen_decode_attention_output_next_q8_scheduling_case(false); - run_qwen_decode_attention_output_next_q8_scheduling_case(true); - run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); - return 0; + + const bool run_vector_cap_experiment = std::getenv("GGML_HRX_RUN_VECTOR_CAP_EXPERIMENT") != nullptr; + const bool run_vector_publish_test = std::getenv("GGML_HRX_RUN_VECTOR_PUBLISH_TEST") != nullptr; + const bool run_vector_postops_test = std::getenv("GGML_HRX_RUN_VECTOR_POSTOPS_TEST") != nullptr; + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + if (!options.runner_control && run_vector_cap_experiment) { + run_vector_output_capacity_experiment(); + return 0; + } + if (!options.runner_control && run_vector_publish_test) { + run_vector_q8_publish_cpu_reference_case(); + return 0; + } + if (!options.runner_control && run_vector_postops_test) { + run_vector_postops_cpu_reference_case(); + return 0; + } + + return suite.run(options, has_device); } diff --git a/tests/testing_suite.h b/tests/testing_suite.h new file mode 100644 index 000000000000..ffa5fa0e66cb --- /dev/null +++ b/tests/testing_suite.h @@ -0,0 +1,548 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#endif + +namespace test_runner { + +struct Case { + std::string name; + std::function run; + bool requires_device = true; +}; + +enum class OutputMode { + Failed, + All, + None, +}; + +struct Config { + std::string suite_name = "test"; + std::string list_flag = "--test-list"; + std::string case_flag = "--test-case"; + std::string filter_flag = "--test-filter"; + std::string jobs_flag = "--test-jobs"; + std::string output_flag = "--test-output"; + std::string env_jobs; + std::string env_filter; + std::string env_output; + size_t max_default_jobs = 4; + + static Config with_prefix(std::string suite_name, + const std::string & flag_prefix, + const std::string & env_prefix, + size_t max_default_jobs = 4) { + Config config; + config.suite_name = std::move(suite_name); + config.list_flag = "--" + flag_prefix + "-test-list"; + config.case_flag = "--" + flag_prefix + "-test-case"; + config.filter_flag = "--" + flag_prefix + "-test-filter"; + config.jobs_flag = "--" + flag_prefix + "-test-jobs"; + config.output_flag = "--" + flag_prefix + "-test-output"; + config.env_jobs = env_prefix + "_JOBS"; + config.env_filter = env_prefix + "_FILTER"; + config.env_output = env_prefix + "_OUTPUT"; + config.max_default_jobs = max_default_jobs; + return config; + } +}; + +struct Options { + std::string executable; + std::string single_case; + std::string filter; + size_t jobs = 0; + bool help = false; + bool list = false; + bool runner_control = false; + OutputMode output = OutputMode::Failed; +}; + +struct Result { + std::string name; + std::string output; + int exit_code = -1; + int signal_code = 0; + int64_t elapsed_ms = 0; + bool passed = false; +}; + +struct Selection { + std::vector cases; + size_t matched = 0; + size_t skipped_device = 0; +}; + +inline void add_case(std::vector & cases, + std::string name, + std::function run, + bool requires_device = true) { + cases.push_back({ std::move(name), std::move(run), requires_device }); +} + +inline size_t default_jobs(const Config & config) { + const unsigned int hardware_jobs = std::thread::hardware_concurrency(); + if (hardware_jobs == 0) { + return 1; + } + return std::min(config.max_default_jobs, hardware_jobs); +} + +inline bool parse_positive_size(const char * value, size_t & out) { + if (value == nullptr || *value == '\0') { + return false; + } + char * end = nullptr; + errno = 0; + const unsigned long parsed = std::strtoul(value, &end, 10); + if (errno != 0 || end == value || *end != '\0' || parsed == 0) { + return false; + } + out = static_cast(parsed); + return true; +} + +inline bool parse_output_mode(const char * value, OutputMode & out) { + if (value == nullptr) { + return true; + } + const std::string mode = value; + if (mode == "failed" || mode == "fail") { + out = OutputMode::Failed; + return true; + } + if (mode == "all") { + out = OutputMode::All; + return true; + } + if (mode == "none") { + out = OutputMode::None; + return true; + } + return false; +} + +inline void print_usage(const Config & config, const char * executable) { + std::fprintf(stderr, + "usage: %s [%s N] [%s REGEX] [%s MODE] [%s]\n" + " %s %s NAME\n" + "\n" + "environment: %s, %s, %s\n" + "output modes: failed, all, none\n", + executable, config.jobs_flag.c_str(), config.filter_flag.c_str(), config.output_flag.c_str(), + config.list_flag.c_str(), executable, config.case_flag.c_str(), config.env_jobs.c_str(), + config.env_filter.c_str(), config.env_output.c_str()); +} + +inline bool parse_options(int argc, char ** argv, const Config & config, Options & options) { + options.executable = argc > 0 ? argv[0] : config.suite_name.c_str(); + if (!config.env_filter.empty()) { + if (const char * env_filter = std::getenv(config.env_filter.c_str())) { + options.filter = env_filter; + options.runner_control = true; + } + } + if (!config.env_jobs.empty()) { + if (const char * env_jobs = std::getenv(config.env_jobs.c_str())) { + if (!parse_positive_size(env_jobs, options.jobs)) { + std::fprintf(stderr, "invalid %s value: %s\n", config.env_jobs.c_str(), env_jobs); + return false; + } + options.runner_control = true; + } + } + if (!config.env_output.empty()) { + if (const char * env_output = std::getenv(config.env_output.c_str())) { + if (!parse_output_mode(env_output, options.output)) { + std::fprintf(stderr, "invalid %s value: %s\n", config.env_output.c_str(), env_output); + return false; + } + options.runner_control = true; + } + } + + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == config.list_flag) { + options.list = true; + options.runner_control = true; + } else if (arg == config.case_flag) { + if (++i >= argc) { + std::fprintf(stderr, "%s requires a case name\n", config.case_flag.c_str()); + return false; + } + options.single_case = argv[i]; + options.runner_control = true; + } else if (arg == config.filter_flag) { + if (++i >= argc) { + std::fprintf(stderr, "%s requires a regex\n", config.filter_flag.c_str()); + return false; + } + options.filter = argv[i]; + options.runner_control = true; + } else if (arg == config.jobs_flag) { + if (++i >= argc || !parse_positive_size(argv[i], options.jobs)) { + std::fprintf(stderr, "%s requires a positive integer\n", config.jobs_flag.c_str()); + return false; + } + options.runner_control = true; + } else if (arg == config.output_flag) { + if (++i >= argc || !parse_output_mode(argv[i], options.output)) { + std::fprintf(stderr, "%s requires failed, all, or none\n", config.output_flag.c_str()); + return false; + } + options.runner_control = true; + } else if (arg == "--help" || arg == "-h") { + options.help = true; + options.runner_control = true; + } else { + std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); + return false; + } + } + + if (options.jobs == 0) { + options.jobs = default_jobs(config); + } + return true; +} + +inline const Case * find_case(const std::vector & cases, const std::string & name) { + for (const Case & test_case : cases) { + if (test_case.name == name) { + return &test_case; + } + } + return nullptr; +} + +inline int run_single_case_in_process(const Config & config, + const std::vector & cases, + const std::string & name, + bool has_device) { + const Case * test_case = find_case(cases, name); + if (test_case == nullptr) { + std::fprintf(stderr, "unknown test case: %s\n", name.c_str()); + return 2; + } + if (test_case->requires_device && !has_device) { + std::fprintf(stderr, "test skipped: no device available for %s\n", config.suite_name.c_str()); + return 0; + } + test_case->run(); + return 0; +} + +inline Result run_single_case_child(const Config & config, const Options & options, const Case & test_case) { + Result result; + result.name = test_case.name; + + const auto start = std::chrono::steady_clock::now(); + +#if defined(_WIN32) + result.output = "process-isolated test runner is not supported on this platform\n"; + result.exit_code = 127; + result.elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + return result; +#else + int pipe_fds[2] = { -1, -1 }; + if (pipe(pipe_fds) != 0) { + result.output = "failed to create output pipe\n"; + result.exit_code = 127; + return result; + } + + std::fflush(nullptr); + const pid_t pid = fork(); + if (pid == 0) { + close(pipe_fds[0]); + dup2(pipe_fds[1], STDOUT_FILENO); + dup2(pipe_fds[1], STDERR_FILENO); + close(pipe_fds[1]); + + std::vector child_argv; + child_argv.push_back(const_cast(options.executable.c_str())); + child_argv.push_back(const_cast(config.case_flag.c_str())); + child_argv.push_back(const_cast(test_case.name.c_str())); + child_argv.push_back(nullptr); + execvp(child_argv[0], child_argv.data()); + std::fprintf(stderr, "failed to exec %s: %s\n", child_argv[0], std::strerror(errno)); + _exit(127); + } + + close(pipe_fds[1]); + if (pid < 0) { + close(pipe_fds[0]); + result.output = "failed to fork child test process\n"; + result.exit_code = 127; + return result; + } + + char buffer[4096]; + for (;;) { + const ssize_t bytes = read(pipe_fds[0], buffer, sizeof(buffer)); + if (bytes > 0) { + result.output.append(buffer, static_cast(bytes)); + continue; + } + if (bytes == 0) { + break; + } + if (errno == EINTR) { + continue; + } + result.output += "failed to read child output\n"; + break; + } + close(pipe_fds[0]); + + int status = 0; + while (waitpid(pid, &status, 0) < 0) { + if (errno == EINTR) { + continue; + } + result.output += "failed to wait for child test process\n"; + result.exit_code = 127; + return result; + } + + if (WIFEXITED(status)) { + result.exit_code = WEXITSTATUS(status); + result.passed = result.exit_code == 0; + } else if (WIFSIGNALED(status)) { + result.signal_code = WTERMSIG(status); + result.passed = false; + } + + result.elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + return result; +#endif +} + +inline bool should_print_output(OutputMode mode, const Result & result) { + if (mode == OutputMode::All) { + return true; + } + if (mode == OutputMode::None) { + return false; + } + return !result.passed; +} + +inline int run_cases_parallel(const Config & config, + const Options & options, + const std::vector & selected_cases) { + if (selected_cases.empty()) { + std::printf("No %s cases selected\n", config.suite_name.c_str()); + return 0; + } + + const size_t jobs = std::min(options.jobs, selected_cases.size()); + std::vector results(selected_cases.size()); + std::atomic next_index(0); + std::atomic finished(0); + std::mutex print_mutex; + const auto suite_start = std::chrono::steady_clock::now(); + + std::printf("Running %zu %s case(s) with %zu worker process(es)\n", selected_cases.size(), + config.suite_name.c_str(), jobs); + + auto worker = [&] { + for (;;) { + const size_t index = next_index.fetch_add(1); + if (index >= selected_cases.size()) { + return; + } + Result result = run_single_case_child(config, options, selected_cases[index]); + const size_t done = finished.fetch_add(1) + 1; + { + std::lock_guard lock(print_mutex); + std::printf("%s %4zu/%zu %-86s %8lld ms\n", result.passed ? "[PASS]" : "[FAIL]", done, + selected_cases.size(), result.name.c_str(), static_cast(result.elapsed_ms)); + if (!result.passed) { + if (result.signal_code != 0) { + std::printf(" terminated by signal %d\n", result.signal_code); + } else { + std::printf(" exit code %d\n", result.exit_code); + } + } + if (should_print_output(options.output, result) && !result.output.empty()) { + std::printf("----- output: %s -----\n%s", result.name.c_str(), result.output.c_str()); + if (result.output.back() != '\n') { + std::printf("\n"); + } + std::printf("----- end output: %s -----\n", result.name.c_str()); + } + std::fflush(stdout); + } + results[index] = std::move(result); + } + }; + + std::vector workers; + workers.reserve(jobs); + for (size_t i = 0; i < jobs; ++i) { + workers.emplace_back(worker); + } + for (std::thread & thread : workers) { + thread.join(); + } + + size_t passed = 0; + for (const Result & result : results) { + if (result.passed) { + ++passed; + } + } + const size_t failed = results.size() - passed; + const int64_t elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - suite_start).count(); + + std::printf("\n%s summary: %zu passed, %zu failed, %zu total, %lld ms\n", config.suite_name.c_str(), passed, failed, + results.size(), static_cast(elapsed_ms)); + if (failed != 0) { + std::printf("Failed cases:\n"); + for (const Result & result : results) { + if (!result.passed) { + std::printf(" %s\n", result.name.c_str()); + } + } + } + return failed == 0 ? 0 : 1; +} + +inline bool select_cases(const std::vector & cases, + const std::string & filter_text, + bool has_device, + bool include_unavailable_device_cases, + Selection & selection) { + try { + std::regex filter; + const bool has_filter = !filter_text.empty(); + if (has_filter) { + filter = std::regex(filter_text); + } + for (const Case & test_case : cases) { + if (has_filter && !std::regex_search(test_case.name, filter)) { + continue; + } + ++selection.matched; + if (!has_device && test_case.requires_device && !include_unavailable_device_cases) { + ++selection.skipped_device; + continue; + } + selection.cases.push_back(test_case); + } + } catch (const std::regex_error & error) { + std::fprintf(stderr, "invalid test filter regex: %s\n", error.what()); + return false; + } + return true; +} + +inline void print_case_list(const std::vector & cases) { + for (const Case & test_case : cases) { + std::printf("[%s] %s\n", test_case.requires_device ? "device" : "host", test_case.name.c_str()); + } +} + +inline bool validate_case_names(const std::vector & cases) { + for (size_t i = 0; i < cases.size(); ++i) { + for (size_t j = i + 1; j < cases.size(); ++j) { + if (cases[i].name == cases[j].name) { + std::fprintf(stderr, "duplicate test case: %s\n", cases[i].name.c_str()); + return false; + } + } + } + return true; +} + +class Suite { + public: + explicit Suite(Config config) : config_(std::move(config)) {} + + template void device_case(std::string name, F && run) { + add_case(cases_, std::move(name), std::function(std::forward(run)), true); + } + + template void host_case(std::string name, F && run) { + add_case(cases_, std::move(name), std::function(std::forward(run)), false); + } + + bool parse_options(int argc, char ** argv, Options & options) const { + return test_runner::parse_options(argc, argv, config_, options); + } + + void print_usage(const char * executable) const { test_runner::print_usage(config_, executable); } + + int run(int argc, char ** argv, bool has_device) const { + Options options; + if (!parse_options(argc, argv, options)) { + print_usage(argc > 0 ? argv[0] : config_.suite_name.c_str()); + return 2; + } + return run(options, has_device); + } + + int run(const Options & options, bool has_device) const { + if (options.help) { + print_usage(options.executable.c_str()); + return 0; + } + if (!validate_case_names(cases_)) { + return 2; + } + if (!options.single_case.empty()) { + return run_single_case_in_process(config_, cases_, options.single_case, has_device); + } + + Selection selection; + if (!select_cases(cases_, options.filter, has_device, options.list, selection)) { + return 2; + } + + if (options.list) { + print_case_list(selection.cases); + return 0; + } + + if (selection.skipped_device != 0) { + std::fprintf(stderr, "test skipped: no device available for %s (%zu device case(s))\n", + config_.suite_name.c_str(), selection.skipped_device); + if (selection.cases.empty()) { + return 0; + } + } + + return run_cases_parallel(config_, options, selection.cases); + } + + private: + Config config_; + std::vector cases_; +}; + +} // namespace test_runner From 1632fe09b3cd4b1548c22c15029691364616af55 Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 15:41:12 -0700 Subject: [PATCH 6/8] Fix HRX flash attention decode row bounds Guard decode split reduction writes by the query row within each KV head instead of the absolute query head index. This prevents inactive rows from writing outputs for the next KV head. Validated with controlled llamae decode A/B and the Lemonade HRX correctness sweep. Assisted-by: Codex --- .../ops/flash_attention_decode_split_f32_f16_wmma.loom | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom index b60cdea28549..a0ca2ecb3afb 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom @@ -650,8 +650,9 @@ template.def<@ggml.flash_attention.decode_split.reduce_completed.cooperative> de %query_row1 = index.add %query_row0, %c1 : index %query_head0 = index.add %query_head_base, %query_row0 : index %query_head1 = index.add %query_head0, %c1 : index - %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + // Inactive rows must not write query heads owned by the next KV head. + %query_head_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index + %query_head_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index %safe_query_row0 = scf.select %query_head_valid0, %query_row0, %c0 : index %safe_query_row1 = scf.select %query_head_valid1, %query_row1, %c0 : index %reducer_active0 = scalar.andi %lane_has_block, %query_head_valid0 : i1 From bfc59f970c0b22d230d62d234498bfba3f5c5cae Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 16:28:36 -0700 Subject: [PATCH 7/8] ggml-hrx: stage common routed mul_mat_id kernels Move common routed F32 mul_mat_id routes onto the staged skinny/tiled kernel family and drop the deprecated common F32 legacy symbols once the staged routes cover them. Keep the Qwen-used F16 routed kernels intact. Assisted-by: Codex --- .../common/dispatch-gated-mul-mat-id.cpp | 12 +- .../common/dispatch-mul-mat-id.cpp | 39 +- .../kernels/loom-libs/manifest.json | 410 +++++++++++++++--- ...inny_input_f32_dequant_accumulate_f32.loom | 302 +++++++++++++ .../matmul_id/skinny_postprocess_f32.loom | 167 +++++++ ...led_input_f32_dequant_accumulate_f32.loom} | 22 +- .../tiled_postprocess_f32.loom} | 16 +- ..._f32_postops_next_rmsnorm_publish_f32.loom | 98 +++++ .../skinny_input_f32_postops_publish_f32.loom | 93 ++++ .../skinny_input_f32_publish_f32.loom | 81 ++++ ...nny_pair_input_f32_swiglu_publish_f32.loom | 405 +++++++++++++++++ ...f32_postops_next_rmsnorm_publish_f32.loom} | 32 +- .../tiled_input_f32_postops_publish_f32.loom} | 28 +- .../tiled_input_f32_publish_f32.loom} | 16 +- ...ed_pair_input_f32_swiglu_publish_f32.loom} | 22 +- tests/hrx-backend-test.cpp | 72 ++- 16 files changed, 1648 insertions(+), 167 deletions(-) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/{mul_mat_id_f32_f32_wmma_core.loom => matmul_id/tiled_input_f32_dequant_accumulate_f32.loom} (96%) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/{mul_mat_id_f32_f32_postops.loom => matmul_id/tiled_postprocess_f32.loom} (88%) create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom => matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom} (76%) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{mul_mat_id_postops_f32_f32_wmma.loom => matmul_id/tiled_input_f32_postops_publish_f32.loom} (75%) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{mul_mat_id_f32_f32_wmma.loom => matmul_id/tiled_input_f32_publish_f32.loom} (78%) rename ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/{mul_mat_id_swiglu_f32_f32_wmma.loom => matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom} (93%) diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp index 722949aaf0f4..901ddeb9e9b1 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp @@ -12,8 +12,10 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kMulMatIdSwiGLUF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_swiglu_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatIdTiledSwiGLUF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnySwiGLUF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32"); struct MulMatIdSwiGLUMatch { const Value * input = nullptr; @@ -146,6 +148,10 @@ static CommonMulMatIdMatch routing_match_for_swiglu(const MulMatIdSwiGLUMatch & return routed; } +static KernelCatalogRef mul_mat_id_swiglu_f32_kernel(const MulMatIdSwiGLUMatch & match) { + return match.token_count <= 5 ? kMulMatIdSkinnySwiGLUF32F32Kernel : kMulMatIdTiledSwiGLUF32F32Kernel; +} + static bool match_mul_mat_id_swiglu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const MulMatIdSwiGLUMatch match = match_mul_mat_id_swiglu(context); if (!match.matched()) { @@ -159,7 +165,7 @@ static bool match_mul_mat_id_swiglu_dispatch(const DispatchMatchContext & contex } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatIdSwiGLUF32F32WmmaKernel); + dispatch.kernel = make_kernel_specialization(mul_mat_id_swiglu_f32_kernel(match)); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_mul_mat_id_to_config_value(match.token_count)); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp index 1bee27ec8a95..5c9822c5aa9e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp @@ -14,12 +14,18 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kMulMatIdF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_f32_f32_wmma"); -static constexpr KernelCatalogRef kMulMatIdPostOpsF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_postops_f32_f32_wmma"); -static constexpr KernelCatalogRef kMulMatIdPostOpsNextRmsNormF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatIdTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdTiledPostOpsF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyPostOpsF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdTiledPostOpsNextRmsNormF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyPostOpsNextRmsNormF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32"); struct MulMatIdPostOpsMatch { CommonMulMatIdMatch root; @@ -55,6 +61,23 @@ static bool is_bias_shape(const Value & value, int64_t output_size) { return true; } +static bool common_mul_mat_id_is_skinny_route(const CommonMulMatIdMatch & match) { + return match.matched() && match.token_count <= 5; +} + +static KernelCatalogRef common_mul_mat_id_f32_kernel(const CommonMulMatIdMatch & match) { + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyF32F32Kernel : kMulMatIdTiledF32F32Kernel; +} + +static KernelCatalogRef common_mul_mat_id_postops_f32_kernel(const CommonMulMatIdMatch & match, bool has_rmsnorm) { + if (has_rmsnorm) { + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyPostOpsNextRmsNormF32F32Kernel : + kMulMatIdTiledPostOpsNextRmsNormF32F32Kernel; + } + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyPostOpsF32F32Kernel : + kMulMatIdTiledPostOpsF32F32Kernel; +} + static MulMatIdPostOpsMatch match_mul_mat_id_postops(const DispatchMatchContext & context) { MulMatIdPostOpsMatch match; const CommonMulMatIdMatch root = common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); @@ -150,7 +173,7 @@ static MulMatIdPostOpsMatch match_mul_mat_id_postops(const DispatchMatchContext } match.root = root; - match.kernel = match.has_rmsnorm ? kMulMatIdPostOpsNextRmsNormF32F32WmmaKernel : kMulMatIdPostOpsF32F32WmmaKernel; + match.kernel = common_mul_mat_id_postops_f32_kernel(root, match.has_rmsnorm); return match; } @@ -169,7 +192,7 @@ static bool match_mul_mat_id_dispatch(const DispatchMatchContext & context, Disp } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatIdF32F32WmmaKernel); + dispatch.kernel = make_kernel_specialization(common_mul_mat_id_f32_kernel(match)); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_mul_mat_id_to_config_value(match.token_count)); diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index c5c468de228a..92420be1208c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -68,15 +68,9 @@ { "path": "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom" }, - { - "path": "ops/mul_mat_id_f32_f32_wmma.loom" - }, { "path": "ops/mul_mat_id_f16_f16_wmma.loom" }, - { - "path": "ops/mul_mat_id_swiglu_f32_f32_wmma.loom" - }, { "path": "ops/mul_mat_id_swiglu_f16_f16_wmma.loom" }, @@ -87,10 +81,40 @@ "path": "motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom" }, { - "path": "ops/mul_mat_id_postops_f32_f32_wmma.loom" + "path": "ops/matmul_id/tiled_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom" + }, + { + "path": "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom" }, { - "path": "ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom" + "path": "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul_id/tiled_postprocess_f32.loom" + }, + { + "path": "motifs/matmul_id/skinny_postprocess_f32.loom" }, { "path": "ops/moe_routing_tables.loom" @@ -221,9 +245,6 @@ { "path": "motifs/mul_mat_quantized_f16_prefill.loom" }, - { - "path": "motifs/mul_mat_id_f32_f32_wmma_core.loom" - }, { "path": "motifs/mul_mat_id_f16_f16_wmma_core.loom" }, @@ -236,9 +257,6 @@ { "path": "motifs/mul_mat_id_q6k_f16_wmma_projection.loom" }, - { - "path": "motifs/mul_mat_id_f32_f32_postops.loom" - }, { "path": "motifs/mul_mat_f32_f32_postops.loom" }, @@ -2561,10 +2579,124 @@ ] }, { - "name": "ggml_mul_mat_id_f32_f32_wmma", + "name": "ggml_mul_mat_id_f16_f16_wmma", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_f32_f32_wmma", - "source": "ops/mul_mat_id_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_id_f16_f16_wmma", + "source": "ops/mul_mat_id_f16_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_id_f16_f16_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_id_f16_f16_wmma_core.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_id_f16_f16_wmma_core.loom" + ] + }, + { + "name": "ggml_mul_mat_id_swiglu_f16_f16_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_swiglu_f16_f16_wmma", + "source": "ops/mul_mat_id_swiglu_f16_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_id_swiglu_f16_f16_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom", + "motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom", + "motifs/mul_mat_id_q4k_f16_wmma_projection.loom", + "motifs/mul_mat_id_q5k_f16_wmma_projection.loom", + "motifs/mul_mat_id_q6k_f16_wmma_projection.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom", + "motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom", + "motifs/mul_mat_id_q4k_f16_wmma_projection.loom", + "motifs/mul_mat_id_q5k_f16_wmma_projection.loom", + "motifs/mul_mat_id_q6k_f16_wmma_projection.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_input_f32_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2595,10 +2727,10 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_f32_f32_wmma.loom" + "ops/matmul_id/tiled_input_f32_publish_f32.loom" ], "library_sources": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2610,7 +2742,7 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2622,10 +2754,10 @@ ] }, { - "name": "ggml_mul_mat_id_f16_f16_wmma", + "name": "ggml_mul_mat_id_skinny_input_f32_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_f16_f16_wmma", - "source": "ops/mul_mat_id_f16_f16_wmma.loom", + "symbol": "ggml_mul_mat_id_skinny_input_f32_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2641,6 +2773,7 @@ "bindings": [ "input", "expert_table", + "partition_table", "weight", "output" ], @@ -2648,27 +2781,44 @@ "read", "read", "read", + "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom" + "ops/matmul_id/skinny_input_f32_publish_f32.loom" ], "library_sources": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_mul_mat_id_swiglu_f32_f32_wmma", + "name": "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_swiglu_f32_f32_wmma", - "source": "ops/mul_mat_id_swiglu_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2685,8 +2835,9 @@ "input", "expert_table", "partition_table", - "gate_weight", - "up_weight", + "weight", + "bias", + "residual_input", "output" ], "binding_access": [ @@ -2695,15 +2846,18 @@ "read", "read", "read", + "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_swiglu_f32_f32_wmma.loom" + "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom" ], "library_sources": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2715,6 +2869,8 @@ ] }, "compile_dependencies": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2726,10 +2882,10 @@ ] }, { - "name": "ggml_mul_mat_id_swiglu_f16_f16_wmma", + "name": "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_swiglu_f16_f16_wmma", - "source": "ops/mul_mat_id_swiglu_f16_f16_wmma.loom", + "symbol": "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2746,8 +2902,9 @@ "input", "expert_table", "partition_table", - "gate_weight", - "up_weight", + "weight", + "bias", + "residual_input", "output" ], "binding_access": [ @@ -2756,20 +2913,18 @@ "read", "read", "read", + "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_swiglu_f16_f16_wmma.loom" + "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom" ], "library_sources": [ - "motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom", - "motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom", - "motifs/mul_mat_id_q4k_f16_wmma_projection.loom", - "motifs/mul_mat_id_q5k_f16_wmma_projection.loom", - "motifs/mul_mat_id_q6k_f16_wmma_projection.loom", + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2781,11 +2936,8 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_id_swiglu_f16_f16_wmma_body.loom", - "motifs/mul_mat_id_swiglu_f16_f16_accumulate.loom", - "motifs/mul_mat_id_q4k_f16_wmma_projection.loom", - "motifs/mul_mat_id_q5k_f16_wmma_projection.loom", - "motifs/mul_mat_id_q6k_f16_wmma_projection.loom", + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2797,10 +2949,10 @@ ] }, { - "name": "ggml_mul_mat_id_postops_f32_f32_wmma", + "name": "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_postops_f32_f32_wmma", - "source": "ops/mul_mat_id_postops_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2820,7 +2972,10 @@ "weight", "bias", "residual_input", - "output" + "residual_output", + "norm_weight", + "normalized_output", + "completion_counters" ], "binding_access": [ "read", @@ -2829,17 +2984,20 @@ "read", "read", "read", + "read_write", + "read", + "read_write", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_postops_f32_f32_wmma.loom" + "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom" ], "library_sources": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", - "motifs/mul_mat_id_f32_f32_postops.loom", + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2851,8 +3009,8 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", - "motifs/mul_mat_id_f32_f32_postops.loom", + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2864,10 +3022,10 @@ ] }, { - "name": "ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", + "name": "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", - "source": "ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -2908,11 +3066,135 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom" + "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32", + "source": "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom" + ], + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32", + "source": "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom" ], "library_sources": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", - "motifs/mul_mat_id_f32_f32_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -2924,8 +3206,6 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_id_f32_f32_wmma_core.loom", - "motifs/mul_mat_id_f32_f32_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..2c59b2d7d059 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,302 @@ +// Shared routed expert WMMA matmul pipeline for GGML MUL_MAT_ID. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul_id.skinny_f32_f32.publish_vector4 +// ggml.matmul_id.skinny_f32_f32.finish_tile +template.decl @ggml.matmul_id.skinny_f32_f32.core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +// Common descriptor packing uses a 9-bit expert ordinal, a 6-bit 32-row +// partition ordinal, and a 5-bit row count minus one. +func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c511_i32 = scalar.constant 511 : i32 + %expert_i32 = scalar.andi %descriptor, %c511_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c9_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c15_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 511)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.matmul_id.skinny_f32_f32.core> device @ggml_mul_mat_id_skinny_f32_f32_core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index + %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %quant_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %partition_table_noalias, %weight_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %weight : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%bounded_token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %bounded_token_count [lt(%assignment_ordinal, %bounded_token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%bounded_token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %weight_k = index.add %k_origin, %load_k : index + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %route0 = index.rem %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_route0 = index.rem %route0, %bounded_input_route_count : index + %input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_route, %input_k] : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 65535)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 65535)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %assignment_token0 = index.div %bounded_assignment0, %bounded_route_count : index + %assignment_route0 = index.rem %bounded_assignment0, %bounded_route_count : index + %assignment_token1 = index.div %bounded_assignment1, %bounded_route_count : index + %assignment_route1 = index.rem %bounded_assignment1, %bounded_route_count : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul_id.skinny_f32_f32.finish_tile>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %channel_tile, %descriptor_ordinal, %bounded_expert, %route_tile_base, %partition_row_count, %expert_table, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom new file mode 100644 index 000000000000..e07d95cd30e1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom @@ -0,0 +1,167 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4> device @ggml_mul_mat_id_skinny_f32_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4> device @ggml_mul_mat_id_skinny_f32_f32_add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%assignment, %channel], %mask, %c0_f32x4 : view<[%row_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.skinny_f32_f32.store_vector4> device @ggml_mul_mat_id_skinny_f32_f32_store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%assignment, %channel], %mask : vector<4xf32>, view<[%row_count]x[%output_size]xf32>, vector<4xi1> + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight> device @ggml_mul_mat_id_skinny_f32_f32_finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %expert_count = index.assume %expert_count0 [range(%expert_count0, 1, 512)] : index + %output_tile_count = index.assume %output_tile_count0 [range(%output_tile_count0, 1, 512)] : index + %maximum_partition_count = index.assume %maximum_partition_count0 [range(%maximum_partition_count0, 1, 2560)] : index + %descriptor_ordinal, %table_partition_count = index.assume %descriptor_ordinal0, %maximum_partition_count [lt(%descriptor_ordinal0, %maximum_partition_count)] : index, index + %expert, %table_expert_count = index.assume %expert0, %expert_count [lt(%expert0, %expert_count)] : index, index + %route_tile_base, %table_token_count0 = index.assume %route_tile_base0, %token_count [lt(%route_tile_base0, %token_count)] : index, index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %counter_scratch_bytes = index.constant 4 : offset + %rms_scratch_bytes = index.constant 512 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %assignment_table_byte_base = index.scale %expert_count, %c4_bytes : index, offset -> offset + %row_count = index.mul %token_count, %route_count : index + %hidden_size_i32 = index.cast %output_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_channel_ordinal_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %expert_table_noalias, %source_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counters_noalias = buffer.assume.noalias %expert_table, %source, %norm_weight, %normalized_output, %completion_counters : buffer, buffer, buffer, buffer, buffer + %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%expert_count]x[%token_count]xi32> + %source_view = buffer.view %source_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %norm_weight_view = buffer.view %norm_weight_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %normalized_output_view = buffer.view %normalized_output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %completion_counter_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%maximum_partition_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %rms_scratch = buffer.alloca align(16) %rms_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %rms_scratch_view = buffer.view %rms_scratch[%c0_offset] : buffer -> view<128xf32> + kernel.barrier scope(workgroup) ordering(release) + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%descriptor_ordinal] {ordering = acq_rel, scope = device} : i32, view<[%maximum_partition_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %is_last_channel_tile = scalar.cmpi eq, %old_counter, %last_channel_ordinal_i32 : i32 + scf.if %is_last_channel_tile { + kernel.barrier scope(workgroup) ordering(acquire) + scf.for %row_offset = [%c0 to %c32 step %c1] { + %valid_row = index.cmp ult, %row_offset, %partition_row_count : index + scf.if %valid_row { + %assignment_ordinal0 = index.add %route_tile_base, %row_offset : index + %assignment_ordinal, %table_token_count = index.assume %assignment_ordinal0, %token_count [lt(%assignment_ordinal0, %token_count)] : index, index + %assignment_i32 = view.load %assignment_view[%expert, %assignment_ordinal] : view<[%expert_count]x[%token_count]xi32> -> i32 + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %assignment_count = index.assume %assignment, %row_count [lt(%assignment, %row_count)] : index, index + %thread_sum = scf.for %channel = [%workitem to %output_size step %c128](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %rms_scratch_view[%subgroup] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c2 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %rms_scratch_view[%lane] : view<128xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %rms_scratch_view[%c0] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %rms_scratch_view[%c0] : view<128xf32> -> f32 + scf.for %channel = [%workitem to %output_size step %c128] { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %weight_value = view.load %norm_weight_view[%channel] : view<[%output_size]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %normalized = scalar.mulf %scaled, %weight_value : f32 + view.store %normalized, %normalized_output_view[%bounded_assignment, %channel] : f32, view<[%row_count]x[%output_size]xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%descriptor_ordinal] {ordering = release, scope = device} : i32, view<[%maximum_partition_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom similarity index 96% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom index 3c74445dae17..359fcf33e146 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_wmma_core.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom @@ -1,13 +1,13 @@ // Shared routed expert WMMA matmul pipeline for GGML MUL_MAT_ID. // // Op wrappers provide publication and tile-finalization hooks by defining: -// ggml.mul_mat_id_f32_f32_wmma.publish_vector4 -// ggml.mul_mat_id_f32_f32_wmma.finish_tile -template.decl @ggml.mul_mat_id_f32_f32_wmma.core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +// ggml.matmul_id.tiled_f32_f32.publish_vector4 +// ggml.matmul_id.tiled_f32_f32.finish_tile +template.decl @ggml.matmul_id.tiled_f32_f32.core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) @@ -43,7 +43,7 @@ func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) - func.return %expert, %route_tile_base, %partition_row_count : index, index, index } -template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32_wmma_core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.core> device @ggml_mul_mat_id_tiled_f32_f32_core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index @@ -268,7 +268,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 scf.if %writes00 { %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> @@ -276,7 +276,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 scf.if %writes01 { %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> @@ -284,7 +284,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 scf.if %writes10 { %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(subgroup) ordering(acq_rel) vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> @@ -292,10 +292,10 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.core> device @ggml_mul_mat_id_f32_f32 scf.if %writes11 { %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - template.apply<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } kernel.barrier scope(workgroup) ordering(acq_rel) - template.apply<@ggml.mul_mat_id_f32_f32_wmma.finish_tile>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %channel_tile, %descriptor_ordinal, %bounded_expert, %route_tile_base, %partition_row_count, %expert_table, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.finish_tile>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %channel_tile, %descriptor_ordinal, %bounded_expert, %route_tile_base, %partition_row_count, %expert_table, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.barrier scope(workgroup) ordering(acq_rel) } template.return diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_postops.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom similarity index 88% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_postops.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom index d3497b0a1895..3514bac420a7 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_id_f32_f32_postops.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom @@ -1,12 +1,12 @@ -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) -template.def<@ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4> device @ggml_mul_mat_id_f32_f32_wmma_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { +template.def<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4> device @ggml_mul_mat_id_tiled_f32_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index %c1 = index.constant 1 : index %c0_offset = index.constant 0 : offset @@ -19,7 +19,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4> device @ggml_mul_ma template.return %sum : vector<4xf32> } -template.def<@ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4> device @ggml_mul_mat_id_f32_f32_wmma_add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { +template.def<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4> device @ggml_mul_mat_id_tiled_f32_f32_add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index @@ -36,7 +36,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4> device @ggml_mu template.return %sum : vector<4xf32> } -template.def<@ggml.mul_mat_id_f32_f32_wmma.store_vector4> device @ggml_mul_mat_id_f32_f32_wmma_store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.store_vector4> device @ggml_mul_mat_id_tiled_f32_f32_store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index @@ -51,7 +51,7 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.store_vector4> device @ggml_mul_mat_i template.return } -template.def<@ggml.mul_mat_id_f32_f32_wmma.finish_routed_rmsnorm_weight> device @ggml_mul_mat_id_f32_f32_wmma_finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight> device @ggml_mul_mat_id_tiled_f32_f32_finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom new file mode 100644 index 000000000000..7a061772f1f5 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom @@ -0,0 +1,98 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.rms_epsilon : f32 + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_postops_next_rmsnorm_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.skinny_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %residual_output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_postops_next_rmsnorm_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.apply<@ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight>(%token_count, %route_count, %output_size, %expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %descriptor_ordinal, %expert, %route_tile_base, %partition_row_count, %expert_table, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %epsilon = config.get @ggml.mul_mat_id_postops.rms_epsilon : f32 + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %epsilon, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom new file mode 100644 index 000000000000..d35eaee7b34e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom @@ -0,0 +1,93 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_postops_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_postops_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.skinny_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_postops_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_postops_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_postops_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %output, %residual_input, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom new file mode 100644 index 000000000000..a0f41b4dd7eb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom @@ -0,0 +1,81 @@ +// Standalone GGML MUL_MAT_ID fallback for routed expert projections. +// +// The kernel consumes a precomputed expert table. Each table row stores compact +// assignment indices encoded as token * route_count + route. Input can either +// be broadcast across all routes or contain one input row per route. +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_f32_f32_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %output_row_count = index.mul %token_count, %route_count : index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%output_size]xf32> + scf.if %publish_word { + %bounded_assignment, %assignment_count = index.assume %assignment, %output_row_count [lt(%assignment, %output_row_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<4xf32>, view<[%output_row_count]x[%output_size]xf32>, vector<4xi1> + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_f32_f32_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id.input_size : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %input_route_count = config.get @ggml.mul_mat_id.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %weight_format = config.get @ggml.mul_mat_id.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom new file mode 100644 index 000000000000..03e040b89253 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom @@ -0,0 +1,405 @@ +// Routed GGML gate/up MUL_MAT_ID fused with SwiGLU for compact expert tables. +template.decl @ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_pair_f32_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_id_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_swiglu.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_swiglu.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c511_i32 = scalar.constant 511 : i32 + %expert_i32 = scalar.andi %descriptor, %c511_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c9_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c15_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 511)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4> device @ggml_mul_mat_id_skinny_pair_f32_swiglu_apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) { + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %negative_gate = vector.subf %c0_f32x4, %gate : vector<4xf32> + %gate_exp = vector.expf %negative_gate : vector<4xf32> + %denominator = vector.addf %c1_f32x4, %gate_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32x4, %denominator : vector<4xf32> + %silu = vector.mulf %gate, %sigmoid : vector<4xf32> + %result = vector.mulf %silu, %up : vector<4xf32> + template.return %result : vector<4xf32> +} + +template.def<@ggml.matmul_id.skinny_pair_f32_swiglu.body> device @ggml_mul_mat_id_skinny_pair_f32_swiglu_f32_f32_wmma_body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index + %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %output_row_count = index.mul %bounded_token_count, %bounded_route_count : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) + %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) + %gate_weight_expert_bytes = index.scale %bounded_output_size, %gate_weight_row_bytes : index, offset -> offset + %up_weight_expert_bytes = index.scale %bounded_output_size, %up_weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%bounded_input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%bounded_token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf32> + %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %gate_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %up_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %gate_weight_stage_view = buffer.view %gate_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %up_weight_stage_view = buffer.view %up_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %gate_result_fragment_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %up_result_fragment_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %gate_result_physical_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %up_result_physical_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %bounded_token_count [lt(%assignment_ordinal, %bounded_token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%bounded_token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %gate_expert_byte_base = index.scale %bounded_expert, %gate_weight_expert_bytes : index, offset -> offset + %up_expert_byte_base = index.scale %bounded_expert, %up_weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%gate_acc00 = %gate_block_acc00 : vector<8xf16>, %gate_acc01 = %gate_block_acc01 : vector<8xf16>, %gate_acc10 = %gate_block_acc10 : vector<8xf16>, %gate_acc11 = %gate_block_acc11 : vector<8xf16>, %up_acc00 = %up_block_acc00 : vector<8xf16>, %up_acc01 = %up_block_acc01 : vector<8xf16>, %up_acc10 = %up_block_acc10 : vector<8xf16>, %up_acc11 = %up_block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %gate_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %gate_weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %gate_expert_byte_base, %channel_byte_add : offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %up_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %up_weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %up_expert_byte_base, %channel_byte_add : offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %gate_weight_values, %gate_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + vector.store %up_weight_values, %up_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_assignment, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %route0 = index.rem %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_route0 = index.rem %route0, %bounded_input_route_count : index + %input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %loaded = vector.load.mask %input_view[%token, %input_route, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_route_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 = scf.for %k_half = [%c0 to %c32 step %c16](%gate_half_acc00 = %gate_acc00 : vector<8xf16>, %gate_half_acc01 = %gate_acc01 : vector<8xf16>, %gate_half_acc10 = %gate_acc10 : vector<8xf16>, %gate_half_acc11 = %gate_acc11 : vector<8xf16>, %up_half_acc00 = %up_acc00 : vector<8xf16>, %up_half_acc01 = %up_acc01 : vector<8xf16>, %up_half_acc10 = %up_acc10 : vector<8xf16>, %up_half_acc11 = %up_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %gate_lhs0 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %gate_lhs1 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs0 = vector.fragment.load %up_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs1 = vector.fragment.load %up_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %gate_half_next00 = vector.mma %gate_lhs0, %rhs0, %gate_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next01 = vector.mma %gate_lhs0, %rhs1, %gate_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next10 = vector.mma %gate_lhs1, %rhs0, %gate_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next11 = vector.mma %gate_lhs1, %rhs1, %gate_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next00 = vector.mma %up_lhs0, %rhs0, %up_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next01 = vector.mma %up_lhs0, %rhs1, %up_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next10 = vector.mma %up_lhs1, %rhs0, %up_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next11 = vector.mma %up_lhs1, %rhs1, %up_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %gate_half_next00, %gate_half_next01, %gate_half_next10, %gate_half_next11, %up_half_next00, %up_half_next01, %up_half_next10, %up_half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 65535)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 65535)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %gate_result00, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result00, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result01, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result01, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result10, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result10, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result11, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result11, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.skinny_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_swiglu.input_size : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_swiglu.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %gate_weight_format = config.get @ggml.mul_mat_id_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_id_swiglu.up_weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_pair_f32_swiglu.body>(%gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom similarity index 76% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom index 51d0c06d1eaa..912fc0c9292c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_next_rmsnorm_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom @@ -1,18 +1,18 @@ -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_routed_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -amdgpu.target @ggml_mul_mat_id_postops_next_rmsnorm_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -34,7 +34,7 @@ config.decl @ggml.mul_mat_id_postops.rms_epsilon : f32 config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] -template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_id_postops_next_rmsnorm_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_postops_next_rmsnorm_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index %c0 = index.constant 0 : index @@ -42,28 +42,28 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat %has_residual = index.cmp ne, %has_residual_value, %c0 : index scf.if %publish_word { %with_bias = scf.if %has_bias -> (vector<4xf32>) { - %biased = template.apply<@ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + %biased = template.apply<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) scf.yield %biased : vector<4xf32> } else { scf.yield %values : vector<4xf32> } %published = scf.if %has_residual -> (vector<4xf32>) { - %added = template.apply<@ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + %added = template.apply<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) scf.yield %added : vector<4xf32> } else { scf.yield %with_bias : vector<4xf32> } - template.apply<@ggml.mul_mat_id_f32_f32_wmma.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %residual_output) : (index, index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %residual_output) : (index, index, index, index, index, vector<4xf32>, buffer) } template.return } -template.def<@ggml.mul_mat_id_f32_f32_wmma.finish_tile> device @ggml_mul_mat_id_postops_next_rmsnorm_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.apply<@ggml.mul_mat_id_f32_f32_wmma.finish_routed_rmsnorm_weight>(%token_count, %route_count, %output_size, %expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %descriptor_ordinal, %expert, %route_tile_base, %partition_row_count, %expert_table, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, buffer, buffer, buffer, buffer, buffer) +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_postops_next_rmsnorm_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.apply<@ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight>(%token_count, %route_count, %output_size, %expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %descriptor_ordinal, %expert, %route_tile_base, %partition_row_count, %expert_table, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, buffer, buffer, buffer, buffer, buffer) template.return } -kernel.def target(@ggml_mul_mat_id_postops_next_rmsnorm_gfx11_wave64) @ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma(%token_count: index) { +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32(%token_count: index) { %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index %route_count = config.get @ggml.mul_mat_id_postops.route_count : index %output_size = config.get @ggml.mul_mat_id_postops.output_size : index @@ -93,6 +93,6 @@ kernel.def target(@ggml_mul_mat_id_postops_next_rmsnorm_gfx11_wave64) @ggml_mul_ %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %partition_ordinal = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_id_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %epsilon, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %epsilon, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom similarity index 75% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom index 31515b97772e..17b2a61a2c79 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_postops_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom @@ -1,16 +1,16 @@ -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -amdgpu.target @ggml_mul_mat_id_postops_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_postops_publish_f32_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -30,7 +30,7 @@ config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(% config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] -template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_id_postops_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_postops_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index %c0 = index.constant 0 : index @@ -38,27 +38,27 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat %has_residual = index.cmp ne, %has_residual_value, %c0 : index scf.if %publish_word { %with_bias = scf.if %has_bias -> (vector<4xf32>) { - %biased = template.apply<@ggml.mul_mat_id_f32_f32_wmma.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + %biased = template.apply<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) scf.yield %biased : vector<4xf32> } else { scf.yield %values : vector<4xf32> } %published = scf.if %has_residual -> (vector<4xf32>) { - %added = template.apply<@ggml.mul_mat_id_f32_f32_wmma.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + %added = template.apply<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) scf.yield %added : vector<4xf32> } else { scf.yield %with_bias : vector<4xf32> } - template.apply<@ggml.mul_mat_id_f32_f32_wmma.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %output) : (index, index, index, index, index, vector<4xf32>, buffer) } template.return } -template.def<@ggml.mul_mat_id_f32_f32_wmma.finish_tile> device @ggml_mul_mat_id_postops_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_postops_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { template.return } -kernel.def target(@ggml_mul_mat_id_postops_gfx11_wave64) @ggml_mul_mat_id_postops_f32_f32_wmma(%token_count: index) { +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_postops_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_postops_publish_f32(%token_count: index) { %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index %route_count = config.get @ggml.mul_mat_id_postops.route_count : index %output_size = config.get @ggml.mul_mat_id_postops.output_size : index @@ -88,6 +88,6 @@ kernel.def target(@ggml_mul_mat_id_postops_gfx11_wave64) @ggml_mul_mat_id_postop %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %partition_ordinal = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_id_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %output, %residual_input, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %output, %residual_input, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom similarity index 78% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom index d329c7fd38b7..45a6e946eddb 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom @@ -3,13 +3,13 @@ // The kernel consumes a precomputed expert table. Each table row stores compact // assignment indices encoded as token * route_count + route. Input can either // be broadcast across all routes or contain one input row per route. -template.decl @ggml.mul_mat_id_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -template.decl @ggml.mul_mat_id_f32_f32_wmma.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) -amdgpu.target @ggml_mul_mat_id_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] @@ -25,7 +25,7 @@ config.decl @ggml.mul_mat_id.weight_format : %value: index where [range(%value, config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] -template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_id_f32_f32_wmma_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_f32_f32_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index @@ -42,11 +42,11 @@ template.def<@ggml.mul_mat_id_f32_f32_wmma.publish_vector4> device @ggml_mul_mat template.return } -template.def<@ggml.mul_mat_id_f32_f32_wmma.finish_tile> device @ggml_mul_mat_id_f32_f32_wmma_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_f32_f32_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { template.return } -kernel.def target(@ggml_mul_mat_id_gfx11_wave64) @ggml_mul_mat_id_f32_f32_wmma(%token_count: index) { +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_publish_f32(%token_count: index) { %expert_count = config.get @ggml.mul_mat_id.expert_count : index %route_count = config.get @ggml.mul_mat_id.route_count : index %output_size = config.get @ggml.mul_mat_id.output_size : index @@ -76,6 +76,6 @@ kernel.def target(@ggml_mul_mat_id_gfx11_wave64) @ggml_mul_mat_id_f32_f32_wmma(% %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %partition_ordinal = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_id_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_swiglu_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom similarity index 93% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_swiglu_f32_f32_wmma.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom index 129ccf24ec26..56a2349ddfab 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_id_swiglu_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom @@ -1,9 +1,9 @@ // Routed GGML gate/up MUL_MAT_ID fused with SwiGLU for compact expert tables. -template.decl @ggml.mul_mat_id_swiglu.apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) +template.decl @ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) -template.decl @ggml.mul_mat_id_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) +template.decl @ggml.matmul_id.tiled_pair_f32_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) -amdgpu.target @ggml_mul_mat_id_swiglu_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.mul_mat_id_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] @@ -53,7 +53,7 @@ func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) - func.return %expert, %route_tile_base, %partition_row_count : index, index, index } -template.def<@ggml.mul_mat_id_swiglu.apply_vector4> device @ggml_mul_mat_id_swiglu_apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) { +template.def<@ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4> device @ggml_mul_mat_id_tiled_pair_f32_swiglu_apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) { %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %c1_f32x4 = vector.constant 1.0 : vector<4xf32> %negative_gate = vector.subf %c0_f32x4, %gate : vector<4xf32> @@ -65,7 +65,7 @@ template.def<@ggml.mul_mat_id_swiglu.apply_vector4> device @ggml_mul_mat_id_swig template.return %result : vector<4xf32> } -template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f32_wmma_body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { +template.def<@ggml.matmul_id.tiled_pair_f32_swiglu.body> device @ggml_mul_mat_id_tiled_pair_f32_swiglu_f32_f32_wmma_body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index @@ -322,7 +322,7 @@ template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f3 %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> - %activated = template.apply<@ggml.mul_mat_id_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %activated = template.apply<@ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> vector.store.mask %activated, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> } @@ -335,7 +335,7 @@ template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f3 %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> - %activated = template.apply<@ggml.mul_mat_id_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %activated = template.apply<@ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> vector.store.mask %activated, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> } @@ -348,7 +348,7 @@ template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f3 %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> - %activated = template.apply<@ggml.mul_mat_id_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %activated = template.apply<@ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> vector.store.mask %activated, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> } @@ -361,7 +361,7 @@ template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f3 %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> - %activated = template.apply<@ggml.mul_mat_id_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %activated = template.apply<@ggml.matmul_id.tiled_pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> vector.store.mask %activated, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> } @@ -370,7 +370,7 @@ template.def<@ggml.mul_mat_id_swiglu.body> device @ggml_mul_mat_id_swiglu_f32_f3 template.return } -kernel.def target(@ggml_mul_mat_id_swiglu_gfx11_wave64) @ggml_mul_mat_id_swiglu_f32_f32_wmma(%token_count: index) { +kernel.def target(@ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32(%token_count: index) { %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index @@ -400,6 +400,6 @@ kernel.def target(@ggml_mul_mat_id_swiglu_gfx11_wave64) @ggml_mul_mat_id_swiglu_ %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %partition_ordinal = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_id_swiglu.body>(%gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.tiled_pair_f32_swiglu.body>(%gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index cf3761120ac8..ebdc592c5dd2 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -7208,6 +7208,25 @@ static std::string common_mul_mat_weight_format(ggml_type type) { } } +static const char * common_mul_mat_id_f32_kernel_name(int64_t token_count) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_publish_f32"; +} + +static const char * common_mul_mat_id_postops_f32_kernel_name(int64_t token_count, bool has_rmsnorm) { + if (has_rmsnorm) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32"; + } + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"; +} + +static const char * common_mul_mat_id_swiglu_f32_kernel_name(int64_t token_count) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32"; +} + static ggml::hrx::DispatchMatch require_common_mul_mat_id_for_graph(ggml::hrx::Graph & graph, ggml::hrx::CommandPlan & plan, std::vector & covered_nodes, @@ -7232,7 +7251,7 @@ static ggml::hrx::DispatchMatch require_common_mul_mat_id_for_graph(ggml::hrx::G REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, output_index, match)); REQUIRE(match.dispatches.size() == 1); const ggml::hrx::Dispatch & dispatch = match.dispatches[0]; - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == common_mul_mat_id_f32_kernel_name(output->ne[2])); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == output->ne[2]); REQUIRE(dispatch.bindings.size() == 5); REQUIRE(dispatch.bindings[0].value == input_value->id); @@ -7270,9 +7289,9 @@ static CommonMulMatIdSwiGLUTensors build_common_mul_mat_id_swiglu_graph(ggml_con ggml_type up_weight_type, ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, int64_t input_size = 256, - int64_t output_size = 128) { + int64_t output_size = 128, + int64_t token_count = 4) { CommonMulMatIdSwiGLUTensors tensors; - constexpr int64_t token_count = 4; constexpr int64_t route_count = 8; constexpr int64_t input_route_count = 1; constexpr int64_t expert_count = 16; @@ -7319,7 +7338,8 @@ static void require_common_mul_mat_id_swiglu_match(ggml_context * REQUIRE(match.metadata.moe_routing_bundles().size() == 1); const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); REQUIRE(dispatch.bindings.size() == 6); require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.input_size", @@ -7430,7 +7450,7 @@ static void run_common_mul_mat_id_swiglu_dispatch_checks() { ggml::hrx::DispatchMatch match; REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, match)); REQUIRE(kernel_name_for_id(match.dispatches.back().kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); } ggml_free(ctx); @@ -7462,9 +7482,9 @@ static CommonMulMatIdPostOpsTensors build_common_mul_mat_id_postops_graph(ggml_c bool include_bias, bool include_residual, bool include_rmsnorm, - ggml_type weight_type = GGML_TYPE_Q4_K) { + ggml_type weight_type = GGML_TYPE_Q4_K, + int64_t token_count = 4) { CommonMulMatIdPostOpsTensors tensors; - constexpr int64_t token_count = 4; constexpr int64_t route_count = 8; constexpr int64_t input_size = 256; constexpr int64_t output_size = 128; @@ -7610,34 +7630,40 @@ static void run_common_mul_mat_id_postops_dispatch_checks() { { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, false, false); - require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, - false, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, false, + false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, false); - require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", false, - true, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), false, true, + false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, false); - require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, - true, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, + false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, false, GGML_TYPE_BF16); - require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, - true, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, + false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, true); require_common_mul_mat_id_postops_match( - ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", false, true, true); + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), false, true, + true); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, true); require_common_mul_mat_id_postops_match( - ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", true, true, true); + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), true, true, + true); } ggml_free(ctx); @@ -8146,7 +8172,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(kernel_name_for_id(gate_up_match.dispatches[1].kernel.kernel_id) == "loom_libs:ggml_moe_build_expert_partition_table"); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[2].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); REQUIRE(gate_up_match.dispatches[2].bindings.size() == 6); REQUIRE(gate_up_match.dispatches[2].bindings[1].value == gate_up_match.transients[0].value); REQUIRE(gate_up_match.dispatches[2].bindings[2].value == gate_up_match.transients[1].value); @@ -8164,7 +8190,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); + common_mul_mat_id_f32_kernel_name(tensors.gate->ne[2])); } { @@ -8178,7 +8204,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } @@ -8193,7 +8219,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } @@ -8207,7 +8233,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); REQUIRE(down_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); } { @@ -8222,7 +8248,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); REQUIRE(down_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); } ggml_free(ctx); From 7fe530b6efdee6c20b354bdaf20fb578c8fbdd82 Mon Sep 17 00:00:00 2001 From: Rob Suderman Date: Thu, 24 Sep 2026 16:56:27 -0700 Subject: [PATCH 8/8] ggml-hrx: factor matmul routes Move common tiled and skinny matmul dispatch onto the refactored Loom routes, prune obsolete generated benchmark artifacts, and add route/correctness coverage for the migrated paths. Validated with the Lemonade and Qwen HRX correctness gate in /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-lemonade-qwen-correctness-20260924T235113Z. Assisted-by: Codex --- .../benchmarks/loom/llama32_3b_f16.loom | 691 -------- .../benchmarks/loom/llama32_3b_f16.pp256.json | 786 --------- .../benchmarks/loom/llama32_3b_f16.pp512.json | 819 --------- .../benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom | 982 ----------- .../loom/qwen3_30b_a3b_q4_k_m.pp256.json | 1073 ------------ .../loom/qwen3_30b_a3b_q4_k_m.pp512.json | 1135 ------------- .../src/ggml-hrx/dispatch/command-program.cpp | 12 +- .../common/dispatch-mul-mat.cpp | 173 +- .../kernels/loom-libs/manifest.json | 897 +++++----- ...inny_input_f32_dequant_accumulate_f32.loom | 1459 +++++++++++++++++ .../motifs/matmul/skinny_postprocess_f32.loom | 31 + .../motifs/matmul/skinny_publish_f32.loom | 14 + .../matmul/skinny_input_f32_publish_f32.loom | 188 +++ .../matmul/tiled_input_f32_publish_f32.loom | 2 +- .../ops/mul_mat_add_f32_f32_wmma.loom | 55 - ...mul_mat_add_next_rmsnorm_f32_f32_wmma.loom | 68 - .../ops/mul_mat_bias_add_f32_f32_wmma.loom | 58 - ...at_bias_add_next_rmsnorm_f32_f32_wmma.loom | 63 - .../ops/mul_mat_bias_f32_f32_wmma.loom | 55 - .../loom-libs/ops/mul_mat_f32_f32_wmma.loom | 82 - .../runtime/command-program-executor.cpp | 226 ++- .../benchmarks/generate-model-benchmarks.py | 11 +- tests/hrx-backend-test.cpp | 159 +- tests/test-hrx-ops.cpp | 115 +- 24 files changed, 2569 insertions(+), 6585 deletions(-) delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.loom delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json delete mode 100644 ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom create mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_wmma.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_f32_f32_wmma.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom delete mode 100644 ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_f32_f32_wmma.loom diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.loom b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.loom deleted file mode 100644 index 3af4e58634f2..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.loom +++ /dev/null @@ -1,691 +0,0 @@ -// Generated by tools/benchmarks/generate-model-benchmarks.py for llama32_3b_f16 scenarios: pp256, pp512, tg8. -// Regenerate from an HRX command program dump rather than editing by hand. - -target.decl @ggml_binary_f32_gfx11_wave64 - -target.decl @ggml_flash_attention_decode_split_gfx11_wave64 - -target.decl @ggml_flash_attention_gfx11_wave64 - -target.decl @ggml_gather_add_gfx11_wave64 - -target.decl @ggml_get_rows_f32_gfx11_wave64 - -target.decl @ggml_mul_mat_add_gfx11_wave64 - -target.decl @ggml_mul_mat_f32_f32_decode_gfx11_wave64 - -target.decl @ggml_mul_mat_gfx11_wave64 - -target.decl @ggml_mul_mat_swiglu_gfx11_wave64 - -target.decl @ggml_rmsnorm_binary_gfx11_wave32 - -target.decl @ggml_rope_f32_gfx11_wave32 - -target.decl @ggml_rope_set_rows_f32_gfx11_wave32 - -target.decl @ggml_set_rows_gfx11_wave64 - -target.decl @llm_attention_k_matmul_rope_set_rows_gfx11_wave64 - -target.decl @llm_attention_q_matmul_rope_gfx11_wave64 - -target.decl @llm_attention_v_matmul_set_rows_gfx11_wave64 - -kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count$0: index) launch(%element_count$1: index, %lhs: buffer, %rhs: buffer, %output: buffer) - -kernel.decl target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count$5: index) launch(%key_value_token_count$6: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) - -kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count$17: index, %key_value_token_count$18: index) launch(%query_token_count$19: index, %key_value_token_count$20: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) - -kernel.decl target(@ggml_gather_add_gfx11_wave64) @ggml_gather_add_f32(%source_token_count$26: index, %output_token_count$27: index, %hidden_size$28: index) launch(%source_token_count$29: index, %output_token_count$30: index, %hidden_size$31: index, %attention: buffer, %residual: buffer, %output_ids: buffer, %output: buffer) - -kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count$36: index, %row_count$37: index, %hidden_size$38: index) launch(%token_count$39: index, %row_count$40: index, %hidden_size$41: index, %token_ids: buffer, %weight: buffer, %output: buffer) - -kernel.decl target(@ggml_mul_mat_add_gfx11_wave64) @ggml_mul_mat_add_f32_f32_wmma(%token_count$45: index) launch(%token_count$46: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) - -kernel.decl target(@ggml_mul_mat_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_f32_f32_decode_wave64(%token_count$51: index, %input_size$52: index, %output_size$53: index) launch(%token_count$54: index, %input_size$55: index, %output_size$56: index, %input: buffer, %weight: buffer, %output: buffer) - -kernel.decl target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_wmma(%token_count$60: index) launch(%token_count$61: index, %input: buffer, %weight: buffer, %output: buffer) - -kernel.decl target(@ggml_mul_mat_swiglu_gfx11_wave64) @ggml_mul_mat_swiglu_f32_f32_wmma(%token_count$65: index) launch(%token_count$66: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count$71: index) launch(%token_count$72: index, %input: buffer, %rhs: buffer, %output: buffer) - -kernel.decl target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count$76: index) launch(%token_count$77: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) - -kernel.decl target(@ggml_rope_set_rows_f32_gfx11_wave32) @ggml_rope_set_rows_f32(%token_count$83: index, %cache_row_count$84: index) launch(%token_count$85: index, %cache_row_count$86: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) - -kernel.decl target(@ggml_set_rows_gfx11_wave64) @ggml_set_rows(%token_count$93: index, %cache_row_count$94: index, %hidden_size$95: index) launch(%token_count$96: index, %cache_row_count$97: index, %hidden_size$98: index, %rows: buffer, %indices: buffer, %cache: buffer) - -kernel.decl target(@llm_attention_k_matmul_rope_set_rows_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_f32_f32_wmma(%token_count$102: index) launch(%token_count$103: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) - -kernel.decl target(@llm_attention_q_matmul_rope_gfx11_wave64) @llm_attention_q_matmul_rope_f32_f32_wmma(%token_count$111: index) launch(%token_count$112: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) - -kernel.decl target(@llm_attention_v_matmul_set_rows_gfx11_wave64) @llm_attention_v_matmul_set_rows_f32_f32_wmma(%token_count$119: index) launch(%token_count$120: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) - -// Scenario: pp256 -check.case public @llama32_3b_f16_pp256_000_ggml_get_rows_f32_case { - %token_count = check.literal value(64) : index - %row_count = check.literal value(128256) : index - %hidden_size = check.literal value(3072) : index - %token_ids = check.generate.fill value(0) : tensor<64xi32> - %weight = check.generate.fill value(0.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(1.0) : tensor<64x3072xf32> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<64xi32>, tensor<128256x3072xf16>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_000_ggml_get_rows_f32_case> @llama32_3b_f16_pp256_000_ggml_get_rows_f32 - -check.case public @llama32_3b_f16_pp256_001_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(2.0) : tensor<64x3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<64x3072xf32>, tensor<3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_001_ggml_rmsnorm_binary_f32_case> @llama32_3b_f16_pp256_001_ggml_rmsnorm_binary_f32 - -check.case public @llama32_3b_f16_pp256_002_llm_attention_q_matmul_rope_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %positions = check.generate.fill value(0) : tensor<64xi32> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %output = check.generate.fill value(0.0) : tensor<64x24x128xf32> - kernel.launch @llm_attention_q_matmul_rope_f32_f32_wmma[%token_count](%token_count, %input, %weight, %positions, %theta, %freq_factors, %output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64xi32>, tensor<64xf32>, tensor<64xf32>, tensor<64x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_002_llm_attention_q_matmul_rope_f32_f32_wmma_case> @llama32_3b_f16_pp256_002_llm_attention_q_matmul_rope_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_003_llm_attention_v_matmul_set_rows_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<1024x3072xf16> - %indices = check.generate.iota offset(0) step(1) period(256) : tensor<64xi64> - %cache = check.generate.fill value(0.0) : tensor<256x1024xf16> - kernel.launch @llm_attention_v_matmul_set_rows_f32_f32_wmma[%token_count](%token_count, %input, %weight, %indices, %cache) : [index](index, tensor<64x3072xf32>, tensor<1024x3072xf16>, tensor<64xi64>, tensor<256x1024xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_003_llm_attention_v_matmul_set_rows_f32_f32_wmma_case> @llama32_3b_f16_pp256_003_llm_attention_v_matmul_set_rows_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<1024x3072xf16> - %positions = check.generate.fill value(0) : tensor<64xi32> - %indices = check.generate.iota offset(0) step(1) period(256) : tensor<64xi64> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %cache = check.generate.fill value(0.0) : tensor<256x8x128xf16> - kernel.launch @llm_attention_k_matmul_rope_set_rows_f32_f32_wmma[%token_count](%token_count, %input, %weight, %positions, %indices, %theta, %freq_factors, %cache) : [index](index, tensor<64x3072xf32>, tensor<1024x3072xf16>, tensor<64xi32>, tensor<64xi64>, tensor<64xf32>, tensor<64xf32>, tensor<256x8x128xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma_case> @llama32_3b_f16_pp256_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_005_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<64x24x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x256xf16> - %output = check.generate.fill value(1.0) : tensor<64x24x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x24x128xf32>, tensor<256x8x128xf16>, tensor<256x8x128xf16>, tensor<64x256xf16>, tensor<64x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_005_ggml_flash_attention_f32_f16_wmma_case> @llama32_3b_f16_pp256_005_ggml_flash_attention_f32_f16_wmma - -check.case public @llama32_3b_f16_pp256_006_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %residual_input = check.generate.fill value(0.25) : tensor<64x3072xf32> - %residual_output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64x3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_006_ggml_mul_mat_add_f32_f32_wmma_case> @llama32_3b_f16_pp256_006_ggml_mul_mat_add_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_007_ggml_mul_mat_swiglu_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0.0) : tensor<64x3072xf32> - %gate_weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %up_weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %output = check.generate.fill value(1.0) : tensor<64x8192xf32> - kernel.launch @ggml_mul_mat_swiglu_f32_f32_wmma[%token_count](%token_count, %input, %gate_weight, %up_weight, %output) : [index](index, tensor<64x3072xf32>, tensor<8192x3072xf16>, tensor<8192x3072xf16>, tensor<64x8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_007_ggml_mul_mat_swiglu_f32_f32_wmma_case> @llama32_3b_f16_pp256_007_ggml_mul_mat_swiglu_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_008_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x8192xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x8192xf16> - %residual_input = check.generate.fill value(0.25) : tensor<64x3072xf32> - %residual_output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<64x8192xf32>, tensor<3072x8192xf16>, tensor<64x3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_008_ggml_mul_mat_add_f32_f32_wmma_case> @llama32_3b_f16_pp256_008_ggml_mul_mat_add_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_009_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_009_ggml_mul_mat_f32_f32_wmma_case> @llama32_3b_f16_pp256_009_ggml_mul_mat_f32_f32_wmma - -check.case public @llama32_3b_f16_pp256_010_ggml_gather_add_f32_case { - %source_token_count = check.literal value(64) : index - %output_token_count = check.literal value(1) : index - %hidden_size = check.literal value(3072) : index - %attention = check.generate.fill value(0.0) : tensor<64x3072xf32> - %residual = check.generate.fill value(0.0) : tensor<64x3072xf32> - %output_ids = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(1.0) : tensor<1x3072xf32> - kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<64x3072xf32>, tensor<64x3072xf32>, tensor<1xi32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_010_ggml_gather_add_f32_case> @llama32_3b_f16_pp256_010_ggml_gather_add_f32 - -check.case public @llama32_3b_f16_pp256_011_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<1x3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x3072xf32>, tensor<3072xf32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_011_ggml_rmsnorm_binary_f32_case> @llama32_3b_f16_pp256_011_ggml_rmsnorm_binary_f32 - -check.case public @llama32_3b_f16_pp256_012_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(8192) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x8192xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<8192x3072xf16>, tensor<1x8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_012_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp256_012_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_pp256_013_ggml_binary_f32_case { - %element_count = check.literal value(8192) : index - %lhs = check.generate.fill value(2.0) : tensor<8192xf32> - %rhs = check.generate.fill value(3.0) : tensor<8192xf32> - %output = check.generate.fill value(0.0) : tensor<8192xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<8192xf32>, tensor<8192xf32>, tensor<8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_013_ggml_binary_f32_case> @llama32_3b_f16_pp256_013_ggml_binary_f32 - -check.case public @llama32_3b_f16_pp256_014_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(8192) : index - %output_size = check.literal value(3072) : index - %input = check.generate.fill value(1.0) : tensor<1x8192xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x8192xf16> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x8192xf32>, tensor<3072x8192xf16>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_014_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp256_014_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_pp256_015_ggml_binary_f32_case { - %element_count = check.literal value(3072) : index - %lhs = check.generate.fill value(2.0) : tensor<3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<3072xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<3072xf32>, tensor<3072xf32>, tensor<3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_015_ggml_binary_f32_case> @llama32_3b_f16_pp256_015_ggml_binary_f32 - -check.case public @llama32_3b_f16_pp256_016_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(128256) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x128256xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<128256x3072xf16>, tensor<1x128256xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp256_016_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp256_016_ggml_mul_mat_f32_f32_decode_wave64 - -// Scenario: pp512 -check.case public @llama32_3b_f16_pp512_000_ggml_get_rows_f32_case { - %token_count = check.literal value(64) : index - %row_count = check.literal value(128256) : index - %hidden_size = check.literal value(3072) : index - %token_ids = check.generate.fill value(0) : tensor<64xi32> - %weight = check.generate.fill value(0.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(1.0) : tensor<64x3072xf32> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<64xi32>, tensor<128256x3072xf16>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_000_ggml_get_rows_f32_case> @llama32_3b_f16_pp512_000_ggml_get_rows_f32 - -check.case public @llama32_3b_f16_pp512_001_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(2.0) : tensor<64x3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<64x3072xf32>, tensor<3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_001_ggml_rmsnorm_binary_f32_case> @llama32_3b_f16_pp512_001_ggml_rmsnorm_binary_f32 - -check.case public @llama32_3b_f16_pp512_002_llm_attention_q_matmul_rope_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %positions = check.generate.fill value(0) : tensor<64xi32> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %output = check.generate.fill value(0.0) : tensor<64x24x128xf32> - kernel.launch @llm_attention_q_matmul_rope_f32_f32_wmma[%token_count](%token_count, %input, %weight, %positions, %theta, %freq_factors, %output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64xi32>, tensor<64xf32>, tensor<64xf32>, tensor<64x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_002_llm_attention_q_matmul_rope_f32_f32_wmma_case> @llama32_3b_f16_pp512_002_llm_attention_q_matmul_rope_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_003_llm_attention_v_matmul_set_rows_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<1024x3072xf16> - %indices = check.generate.iota offset(0) step(1) period(512) : tensor<64xi64> - %cache = check.generate.fill value(0.0) : tensor<512x1024xf16> - kernel.launch @llm_attention_v_matmul_set_rows_f32_f32_wmma[%token_count](%token_count, %input, %weight, %indices, %cache) : [index](index, tensor<64x3072xf32>, tensor<1024x3072xf16>, tensor<64xi64>, tensor<512x1024xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_003_llm_attention_v_matmul_set_rows_f32_f32_wmma_case> @llama32_3b_f16_pp512_003_llm_attention_v_matmul_set_rows_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<1024x3072xf16> - %positions = check.generate.fill value(0) : tensor<64xi32> - %indices = check.generate.iota offset(0) step(1) period(512) : tensor<64xi64> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %cache = check.generate.fill value(0.0) : tensor<512x8x128xf16> - kernel.launch @llm_attention_k_matmul_rope_set_rows_f32_f32_wmma[%token_count](%token_count, %input, %weight, %positions, %indices, %theta, %freq_factors, %cache) : [index](index, tensor<64x3072xf32>, tensor<1024x3072xf16>, tensor<64xi32>, tensor<64xi64>, tensor<64xf32>, tensor<64xf32>, tensor<512x8x128xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma_case> @llama32_3b_f16_pp512_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_005_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<64x24x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x256xf16> - %output = check.generate.fill value(1.0) : tensor<64x24x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x24x128xf32>, tensor<256x8x128xf16>, tensor<256x8x128xf16>, tensor<64x256xf16>, tensor<64x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_005_ggml_flash_attention_f32_f16_wmma_case> @llama32_3b_f16_pp512_005_ggml_flash_attention_f32_f16_wmma - -check.case public @llama32_3b_f16_pp512_006_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %residual_input = check.generate.fill value(0.25) : tensor<64x3072xf32> - %residual_output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64x3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_006_ggml_mul_mat_add_f32_f32_wmma_case> @llama32_3b_f16_pp512_006_ggml_mul_mat_add_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_007_ggml_mul_mat_swiglu_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0.0) : tensor<64x3072xf32> - %gate_weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %up_weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %output = check.generate.fill value(1.0) : tensor<64x8192xf32> - kernel.launch @ggml_mul_mat_swiglu_f32_f32_wmma[%token_count](%token_count, %input, %gate_weight, %up_weight, %output) : [index](index, tensor<64x3072xf32>, tensor<8192x3072xf16>, tensor<8192x3072xf16>, tensor<64x8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_007_ggml_mul_mat_swiglu_f32_f32_wmma_case> @llama32_3b_f16_pp512_007_ggml_mul_mat_swiglu_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_008_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x8192xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x8192xf16> - %residual_input = check.generate.fill value(0.25) : tensor<64x3072xf32> - %residual_output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<64x8192xf32>, tensor<3072x8192xf16>, tensor<64x3072xf32>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_008_ggml_mul_mat_add_f32_f32_wmma_case> @llama32_3b_f16_pp512_008_ggml_mul_mat_add_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_009_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(1.0) : tensor<64x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %output = check.generate.fill value(0.0) : tensor<64x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<64x3072xf32>, tensor<3072x3072xf16>, tensor<64x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_009_ggml_mul_mat_f32_f32_wmma_case> @llama32_3b_f16_pp512_009_ggml_mul_mat_f32_f32_wmma - -check.case public @llama32_3b_f16_pp512_010_ggml_gather_add_f32_case { - %source_token_count = check.literal value(64) : index - %output_token_count = check.literal value(1) : index - %hidden_size = check.literal value(3072) : index - %attention = check.generate.fill value(0.0) : tensor<64x3072xf32> - %residual = check.generate.fill value(0.0) : tensor<64x3072xf32> - %output_ids = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(1.0) : tensor<1x3072xf32> - kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<64x3072xf32>, tensor<64x3072xf32>, tensor<1xi32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_010_ggml_gather_add_f32_case> @llama32_3b_f16_pp512_010_ggml_gather_add_f32 - -check.case public @llama32_3b_f16_pp512_011_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<1x3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x3072xf32>, tensor<3072xf32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_011_ggml_rmsnorm_binary_f32_case> @llama32_3b_f16_pp512_011_ggml_rmsnorm_binary_f32 - -check.case public @llama32_3b_f16_pp512_012_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(8192) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x8192xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<8192x3072xf16>, tensor<1x8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_012_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp512_012_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_pp512_013_ggml_binary_f32_case { - %element_count = check.literal value(8192) : index - %lhs = check.generate.fill value(2.0) : tensor<8192xf32> - %rhs = check.generate.fill value(3.0) : tensor<8192xf32> - %output = check.generate.fill value(0.0) : tensor<8192xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<8192xf32>, tensor<8192xf32>, tensor<8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_013_ggml_binary_f32_case> @llama32_3b_f16_pp512_013_ggml_binary_f32 - -check.case public @llama32_3b_f16_pp512_014_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(8192) : index - %output_size = check.literal value(3072) : index - %input = check.generate.fill value(1.0) : tensor<1x8192xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x8192xf16> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x8192xf32>, tensor<3072x8192xf16>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_014_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp512_014_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_pp512_015_ggml_binary_f32_case { - %element_count = check.literal value(3072) : index - %lhs = check.generate.fill value(2.0) : tensor<3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<3072xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<3072xf32>, tensor<3072xf32>, tensor<3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_015_ggml_binary_f32_case> @llama32_3b_f16_pp512_015_ggml_binary_f32 - -check.case public @llama32_3b_f16_pp512_016_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(128256) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x128256xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<128256x3072xf16>, tensor<1x128256xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_016_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_pp512_016_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_pp512_017_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(512) : index - %query = check.generate.fill value(0.0) : tensor<64x24x128xf32> - %key = check.generate.fill value(0.0) : tensor<512x8x128xf16> - %value = check.generate.fill value(0.0) : tensor<512x8x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x512xf16> - %output = check.generate.fill value(1.0) : tensor<64x24x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x24x128xf32>, tensor<512x8x128xf16>, tensor<512x8x128xf16>, tensor<64x512xf16>, tensor<64x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_pp512_017_ggml_flash_attention_f32_f16_wmma_case> @llama32_3b_f16_pp512_017_ggml_flash_attention_f32_f16_wmma - -// Scenario: tg8 -check.case public @llama32_3b_f16_tg8_000_ggml_get_rows_f32_case { - %token_count = check.literal value(1) : index - %row_count = check.literal value(128256) : index - %hidden_size = check.literal value(3072) : index - %token_ids = check.generate.fill value(0) : tensor<1xi32> - %weight = check.generate.fill value(0.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(1.0) : tensor<1x3072xf32> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<128256x3072xf16>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_000_ggml_get_rows_f32_case> @llama32_3b_f16_tg8_000_ggml_get_rows_f32 - -check.case public @llama32_3b_f16_tg8_001_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(2.0) : tensor<1x3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x3072xf32>, tensor<3072xf32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_001_ggml_rmsnorm_binary_f32_case> @llama32_3b_f16_tg8_001_ggml_rmsnorm_binary_f32 - -check.case public @llama32_3b_f16_tg8_002_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(3072) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<3072x3072xf16>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_002_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_tg8_002_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_tg8_003_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(1024) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<1024x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x1024xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<1024x3072xf16>, tensor<1x1024xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_003_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_tg8_003_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_tg8_004_ggml_rope_f32_case { - %token_count = check.literal value(1) : index - %positions = check.generate.fill value(0) : tensor<1xi32> - %input = check.generate.fill value(0.0) : tensor<1x24x128xf32> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %output = check.generate.fill value(1.0) : tensor<1x24x128xf32> - kernel.launch @ggml_rope_f32[%token_count](%token_count, %positions, %input, %theta, %freq_factors, %output) : [index](index, tensor<1xi32>, tensor<1x24x128xf32>, tensor<64xf32>, tensor<64xf32>, tensor<1x24x128xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_004_ggml_rope_f32_case> @llama32_3b_f16_tg8_004_ggml_rope_f32 - -check.case public @llama32_3b_f16_tg8_005_ggml_rope_set_rows_f32_case { - %token_count = check.literal value(1) : index - %cache_row_count = check.literal value(256) : index - %positions = check.generate.fill value(0) : tensor<1xi32> - %indices = check.generate.iota offset(0) step(1) period(256) : tensor<1xi64> - %input = check.generate.fill value(0.0) : tensor<1x8x128xf32> - %theta = check.generate.fill value(0.0) : tensor<64xf32> - %freq_factors = check.generate.fill value(0.0) : tensor<64xf32> - %cache = check.generate.fill value(0.0) : tensor<256x8x128xf16> - kernel.launch @ggml_rope_set_rows_f32[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %indices, %input, %theta, %freq_factors, %cache) : [index, index](index, index, tensor<1xi32>, tensor<1xi64>, tensor<1x8x128xf32>, tensor<64xf32>, tensor<64xf32>, tensor<256x8x128xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_005_ggml_rope_set_rows_f32_case> @llama32_3b_f16_tg8_005_ggml_rope_set_rows_f32 - -check.case public @llama32_3b_f16_tg8_006_ggml_set_rows_case { - %token_count = check.literal value(1) : index - %cache_row_count = check.literal value(256) : index - %hidden_size = check.literal value(1024) : index - %rows = check.generate.fill value(0.0) : tensor<1x1024xf32> - %indices = check.generate.iota offset(0) step(1) period(256) : tensor<1xi64> - %cache = check.generate.fill value(0.0) : tensor<256x1024xf16> - kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<1x1024xf32>, tensor<1xi64>, tensor<256x1024xf16>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_006_ggml_set_rows_case> @llama32_3b_f16_tg8_006_ggml_set_rows - -check.case public @llama32_3b_f16_tg8_007_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8_case { - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<24x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x8x128xf16> - %mask = check.generate.fill value(0.0) : tensor<256xf16> - %partial_max = check.generate.fill value(0.0) : tensor<8x4x16xf32> - %partial_sum = check.generate.fill value(0.0) : tensor<8x4x16xf32> - %partial_output = check.generate.fill value(0.0) : tensor<8x4x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<8xi32> - %output = check.generate.fill value(1.0) : tensor<24x128xf32> - %next_q8_output = check.generate.fill value(0) : tensor<3456xi8> - kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %next_q8_output) : [index](index, tensor<24x128xf32>, tensor<256x8x128xf16>, tensor<256x8x128xf16>, tensor<256xf16>, tensor<8x4x16xf32>, tensor<8x4x16xf32>, tensor<8x4x16x128xf16>, tensor<8xi32>, tensor<24x128xf32>, tensor<3456xi8>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_007_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8_case> @llama32_3b_f16_tg8_007_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8 - -check.case public @llama32_3b_f16_tg8_008_ggml_binary_f32_case { - %element_count = check.literal value(3072) : index - %lhs = check.generate.fill value(2.0) : tensor<3072xf32> - %rhs = check.generate.fill value(3.0) : tensor<3072xf32> - %output = check.generate.fill value(0.0) : tensor<3072xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<3072xf32>, tensor<3072xf32>, tensor<3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_008_ggml_binary_f32_case> @llama32_3b_f16_tg8_008_ggml_binary_f32 - -check.case public @llama32_3b_f16_tg8_009_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(8192) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<8192x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x8192xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<8192x3072xf16>, tensor<1x8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_009_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_tg8_009_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_tg8_010_ggml_binary_f32_case { - %element_count = check.literal value(8192) : index - %lhs = check.generate.fill value(2.0) : tensor<8192xf32> - %rhs = check.generate.fill value(3.0) : tensor<8192xf32> - %output = check.generate.fill value(0.0) : tensor<8192xf32> - kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<8192xf32>, tensor<8192xf32>, tensor<8192xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_010_ggml_binary_f32_case> @llama32_3b_f16_tg8_010_ggml_binary_f32 - -check.case public @llama32_3b_f16_tg8_011_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(8192) : index - %output_size = check.literal value(3072) : index - %input = check.generate.fill value(1.0) : tensor<1x8192xf32> - %weight = check.generate.fill value(1.0) : tensor<3072x8192xf16> - %output = check.generate.fill value(0.0) : tensor<1x3072xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x8192xf32>, tensor<3072x8192xf16>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_011_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_tg8_011_ggml_mul_mat_f32_f32_decode_wave64 - -check.case public @llama32_3b_f16_tg8_012_ggml_gather_add_f32_case { - %source_token_count = check.literal value(1) : index - %output_token_count = check.literal value(1) : index - %hidden_size = check.literal value(3072) : index - %attention = check.generate.fill value(0.0) : tensor<1x3072xf32> - %residual = check.generate.fill value(0.0) : tensor<1x3072xf32> - %output_ids = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(1.0) : tensor<1x3072xf32> - kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<1x3072xf32>, tensor<1xi32>, tensor<1x3072xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_012_ggml_gather_add_f32_case> @llama32_3b_f16_tg8_012_ggml_gather_add_f32 - -check.case public @llama32_3b_f16_tg8_013_ggml_mul_mat_f32_f32_decode_wave64_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(3072) : index - %output_size = check.literal value(128256) : index - %input = check.generate.fill value(1.0) : tensor<1x3072xf32> - %weight = check.generate.fill value(1.0) : tensor<128256x3072xf16> - %output = check.generate.fill value(0.0) : tensor<1x128256xf32> - kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<1x3072xf32>, tensor<128256x3072xf16>, tensor<1x128256xf32>) - check.return -} - -check.benchmark<@llama32_3b_f16_tg8_013_ggml_mul_mat_f32_f32_decode_wave64_case> @llama32_3b_f16_tg8_013_ggml_mul_mat_f32_f32_decode_wave64 diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json deleted file mode 100644 index 18306b93d7e2..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp256.json +++ /dev/null @@ -1,786 +0,0 @@ -{ - "command_count": 259, - "dispatch_count": 17, - "dispatches": [ - { - "benchmark": "@llama32_3b_f16_pp256_000_ggml_get_rows_f32", - "compile_parameters": { - "ggml.get_rows_f32.hidden_capacity": "3072", - "ggml.get_rows_f32.token_capacity": "64", - "ggml.get_rows_f32.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "hidden_size": 3072, - "row_count": 128256, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_get_rows_f32", - "library_sources": [ - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/rows/get_rows_f32.loom" - ], - "sources": [ - "ops/rows/get_rows_f32.loom", - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_get_rows_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "row_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_001_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "3072", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999975e-06" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 55, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_002_llm_attention_q_matmul_rope_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.head_count": "24", - "llm.attention_qkv.head_size": "128", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "3072", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_q_matmul_rope_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_003_llm_attention_v_matmul_set_rows_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.cache_output_format": "16", - "llm.attention_qkv.cache_row_count": "256", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "1024", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_v_matmul_set_rows_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.cache_output_format": "16", - "llm.attention_qkv.cache_row_count": "256", - "llm.attention_qkv.head_count": "8", - "llm.attention_qkv.head_size": "128", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "1024", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_005_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "8", - "ggml.flash_attention.query_head_count": "24" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "key_value_token_count": 256, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_006_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "3072", - "ggml.mul_mat_postops.output_size": "3072", - "ggml.mul_mat_postops.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 27, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_007_ggml_mul_mat_swiglu_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_swiglu.gate_weight_format": "16", - "ggml.mul_mat_swiglu.input_size": "3072", - "ggml.mul_mat_swiglu.op": "4", - "ggml.mul_mat_swiglu.output_size": "8192", - "ggml.mul_mat_swiglu.up_weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 27, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_swiglu_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_008_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "8192", - "ggml.mul_mat_postops.output_size": "3072", - "ggml.mul_mat_postops.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 27, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_009_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "3072", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "3072", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_010_ggml_gather_add_f32", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/hrx", - "count": 1, - "integer_parameters": { - "hidden_size": 3072, - "output_token_count": 1, - "source_token_count": 64 - }, - "kernel": "hrx:ggml_gather_add_f32", - "library_sources": [], - "primary_sources": [ - "gather_add_f32.loom" - ], - "sources": [ - "gather_add_f32.loom" - ], - "symbol": "ggml_gather_add_f32", - "workload_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_011_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "3072", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999975e-06" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "token_count": 1 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_012_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "8192", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "input_size": 3072, - "output_size": 8192, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_013_ggml_binary_f32", - "compile_parameters": { - "ggml.binary_f32.op": "4" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "element_count": 8192 - }, - "kernel": "loom_libs:ggml_binary_f32", - "library_sources": [ - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/binary_f32.loom" - ], - "sources": [ - "ops/binary_f32.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_binary_f32", - "workload_parameters": [ - { - "name": "element_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_014_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "3072", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "input_size": 8192, - "output_size": 3072, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_015_ggml_binary_f32", - "compile_parameters": { - "ggml.binary_f32.op": "0" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "element_count": 3072 - }, - "kernel": "loom_libs:ggml_binary_f32", - "library_sources": [ - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/binary_f32.loom" - ], - "sources": [ - "ops/binary_f32.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_binary_f32", - "workload_parameters": [ - { - "name": "element_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp256_016_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "128256", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "input_size": 3072, - "output_size": 128256, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - } - ], - "generated_count": 17, - "kernel_counts": { - "hrx:ggml_gather_add_f32": 1, - "loom_libs:ggml_binary_f32": 2, - "loom_libs:ggml_flash_attention_f32_f16_wmma": 28, - "loom_libs:ggml_get_rows_f32": 1, - "loom_libs:ggml_mul_mat_add_f32_f32_wmma": 54, - "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": 4, - "loom_libs:ggml_mul_mat_f32_f32_wmma": 1, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma": 27, - "loom_libs:ggml_rmsnorm_binary_f32": 57, - "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma": 28, - "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma": 28, - "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma": 28 - }, - "loom_source": "benchmarks/loom/llama32_3b_f16.loom", - "model": "llama32_3b_f16", - "scenario": "pp256", - "schema": "ggml-hrx-model-loom-benchmarks-v2" -} diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json deleted file mode 100644 index 9752830da780..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.pp512.json +++ /dev/null @@ -1,819 +0,0 @@ -{ - "command_count": 518, - "dispatch_count": 18, - "dispatches": [ - { - "benchmark": "@llama32_3b_f16_pp512_000_ggml_get_rows_f32", - "compile_parameters": { - "ggml.get_rows_f32.hidden_capacity": "3072", - "ggml.get_rows_f32.token_capacity": "64", - "ggml.get_rows_f32.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "hidden_size": 3072, - "row_count": 128256, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_get_rows_f32", - "library_sources": [ - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/rows/get_rows_f32.loom" - ], - "sources": [ - "ops/rows/get_rows_f32.loom", - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_get_rows_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "row_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_001_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "3072", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999975e-06" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 110, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_002_llm_attention_q_matmul_rope_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.head_count": "24", - "llm.attention_qkv.head_size": "128", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "3072", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 56, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_q_matmul_rope_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_q_matmul_rope_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_003_llm_attention_v_matmul_set_rows_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.cache_output_format": "16", - "llm.attention_qkv.cache_row_count": "512", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "1024", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 56, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_v_matmul_set_rows_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_v_matmul_set_rows_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_004_llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "compile_parameters": { - "ggml.workload.token_capacity": "64", - "llm.attention_qkv.cache_output_format": "16", - "llm.attention_qkv.cache_row_count": "512", - "llm.attention_qkv.head_count": "8", - "llm.attention_qkv.head_size": "128", - "llm.attention_qkv.input_size": "3072", - "llm.attention_qkv.output_size": "1024", - "llm.attention_qkv.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 56, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom" - ], - "sources": [ - "ops/llm_attention_k_matmul_rope_set_rows_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/llm_attention_qkv_matmul_postops.loom", - "motifs/rope/rope_f32.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_005_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "8", - "ggml.flash_attention.query_head_count": "24" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "key_value_token_count": 256, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_006_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "3072", - "ggml.mul_mat_postops.output_size": "3072", - "ggml.mul_mat_postops.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 54, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_007_ggml_mul_mat_swiglu_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_swiglu.gate_weight_format": "16", - "ggml.mul_mat_swiglu.input_size": "3072", - "ggml.mul_mat_swiglu.op": "4", - "ggml.mul_mat_swiglu.output_size": "8192", - "ggml.mul_mat_swiglu.up_weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 54, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_swiglu_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_008_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "8192", - "ggml.mul_mat_postops.output_size": "3072", - "ggml.mul_mat_postops.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 54, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_009_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "3072", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "3072", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "16", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_010_ggml_gather_add_f32", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/hrx", - "count": 2, - "integer_parameters": { - "hidden_size": 3072, - "output_token_count": 1, - "source_token_count": 64 - }, - "kernel": "hrx:ggml_gather_add_f32", - "library_sources": [], - "primary_sources": [ - "gather_add_f32.loom" - ], - "sources": [ - "gather_add_f32.loom" - ], - "symbol": "ggml_gather_add_f32", - "workload_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_011_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "3072", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999975e-06" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 4, - "integer_parameters": { - "token_count": 1 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_012_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "8192", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 4, - "integer_parameters": { - "input_size": 3072, - "output_size": 8192, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_013_ggml_binary_f32", - "compile_parameters": { - "ggml.binary_f32.op": "4" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "element_count": 8192 - }, - "kernel": "loom_libs:ggml_binary_f32", - "library_sources": [ - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/binary_f32.loom" - ], - "sources": [ - "ops/binary_f32.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_binary_f32", - "workload_parameters": [ - { - "name": "element_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_014_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "3072", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "input_size": 8192, - "output_size": 3072, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_015_ggml_binary_f32", - "compile_parameters": { - "ggml.binary_f32.op": "0" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "element_count": 3072 - }, - "kernel": "loom_libs:ggml_binary_f32", - "library_sources": [ - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/binary_f32.loom" - ], - "sources": [ - "ops/binary_f32.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_binary_f32", - "workload_parameters": [ - { - "name": "element_count", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_016_ggml_mul_mat_f32_f32_decode_wave64", - "compile_parameters": { - "ggml.mul_mat_f32_f32_decode.output_capacity": "128256", - "ggml.mul_mat_f32_f32_decode.token_capacity": "1", - "ggml.mul_mat_f32_f32_decode.weight_format": "16" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "input_size": 3072, - "output_size": 128256, - "token_count": 1 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@llama32_3b_f16_pp512_017_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "8", - "ggml.flash_attention.query_head_count": "24" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 28, - "integer_parameters": { - "key_value_token_count": 512, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - } - ], - "generated_count": 18, - "kernel_counts": { - "hrx:ggml_gather_add_f32": 2, - "loom_libs:ggml_binary_f32": 4, - "loom_libs:ggml_flash_attention_f32_f16_wmma": 56, - "loom_libs:ggml_get_rows_f32": 2, - "loom_libs:ggml_mul_mat_add_f32_f32_wmma": 108, - "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": 8, - "loom_libs:ggml_mul_mat_f32_f32_wmma": 2, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma": 54, - "loom_libs:ggml_rmsnorm_binary_f32": 114, - "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma": 56, - "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma": 56, - "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma": 56 - }, - "loom_source": "benchmarks/loom/llama32_3b_f16.loom", - "model": "llama32_3b_f16", - "scenario": "pp512", - "schema": "ggml-hrx-model-loom-benchmarks-v2" -} diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom deleted file mode 100644 index e78d397230f1..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom +++ /dev/null @@ -1,982 +0,0 @@ -// Generated by tools/benchmarks/generate-model-benchmarks.py for qwen3_30b_a3b_q4_k_m scenarios: pp256, pp512, tg8. -// Regenerate from an HRX command program dump rather than editing by hand. - -target.decl @ggml_flash_attention_decode_split_gfx11_wave64 -target.decl @ggml_flash_attention_gfx11_wave64 -target.decl @ggml_gather_add_gfx11_wave64 -target.decl @ggml_get_rows_f32_gfx11_wave64 -target.decl @ggml_moe_routing_gfx11_wave32 -target.decl @ggml_mul_mat_add_gfx11_wave64 -target.decl @ggml_mul_mat_gfx11_wave64 -target.decl @ggml_mul_mat_id_f16_f16_gfx11_wave64 -target.decl @ggml_rmsnorm_binary_gfx11_wave32 -target.decl @qwen3_moe_attention_postprocess_gfx11_wave32 -target.decl @qwen3_moe_attention_prepare_gfx11_wave32 -target.decl @qwen3_moe_attention_qkv_postprocess_gfx11_wave32 -target.decl @qwen3_moe_gfx11_wave32 -target.decl @qwen3_moe_routed_down_next_norm_gfx11_wave32 -target.decl @qwen3_moe_routed_down_q6k_gfx11_wave64 -target.decl @qwen3_moe_router_fused_gfx11_wave64 -target.decl @qwen3_moe_router_gfx11_wave64 -target.decl @qwen3_moe_router_projection_gfx11_wave32 -target.decl @qwen_attention_metadata_gfx11_wave64 -target.decl @qwen_attention_state_gfx11_wave64 - -kernel.decl target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) -kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) -kernel.decl target(@ggml_gather_add_gfx11_wave64) @ggml_gather_add_f32(%source_token_count: index, %output_token_count: index, %hidden_size: index) launch(%source_token_count: index, %output_token_count: index, %hidden_size: index, %attention: buffer, %residual: buffer, %output_ids: buffer, %output: buffer) -kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) -kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) -kernel.decl target(@ggml_moe_routing_gfx11_wave32) @ggml_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) -kernel.decl target(@ggml_moe_routing_gfx11_wave32) @ggml_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) -kernel.decl target(@ggml_mul_mat_add_gfx11_wave64) @ggml_mul_mat_add_f32_f32_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) -kernel.decl target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) -kernel.decl target(@ggml_mul_mat_id_f16_f16_gfx11_wave64) @ggml_mul_mat_id_f16_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) -kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) -kernel.decl target(@qwen3_moe_attention_postprocess_gfx11_wave32) @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) launch(%token_count: index, %cache_row_count: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) -kernel.decl target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode(%token_count: index, %cache_row_count: index) launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) -kernel.decl target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) -kernel.decl target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) -kernel.decl target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) -kernel.decl target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) -kernel.decl target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32(%token_count: index) launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) -kernel.decl target(@qwen3_moe_gfx11_wave32) @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) -kernel.decl target(@qwen3_moe_router_projection_gfx11_wave32) @qwen3_moe_router_projection_f32_four_row_wave32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) -kernel.decl target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projection_top8_fused_decode_f32(%token_count: index, %route_id_stride: index) launch(%token_count: index, %route_id_stride: index, %input: buffer, %weight: buffer, %logits: buffer, %completion_counter: buffer, %route_ids: buffer, %route_weights: buffer) -kernel.decl target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) launch(%token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) -kernel.decl target(@qwen_attention_state_gfx11_wave64) @qwen_attention_context_base_capture() launch(%positions: buffer, %control: buffer) -kernel.decl target(@qwen_attention_metadata_gfx11_wave64) @qwen_attention_metadata(%token_count: index, %context_capacity: index) launch(%token_count: index, %context_capacity: index, %control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %attention_mask: buffer) - - -// Scenario: pp256 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_000_qwen_attention_context_base_capture_case { - %positions = check.generate.fill value(0) : tensor<256xi8> - %control = check.generate.fill value(0) : tensor<4xi8> - kernel.launch @qwen_attention_context_base_capture[](%positions, %control) : [](tensor<256xi8>, tensor<4xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_000_qwen_attention_context_base_capture_case> @qwen3_30b_a3b_q4_k_m_pp256_000_qwen_attention_context_base_capture - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_001_qwen_attention_metadata_case { - %token_count = check.literal value(64) : index - %context_capacity = check.literal value(256) : index - %control = check.generate.fill value(0) : tensor<4xi8> - %positions = check.generate.fill value(0) : tensor<256xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %attention_mask = check.generate.fill value(0) : tensor<32768xi8> - kernel.launch @qwen_attention_metadata[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<4xi8>, tensor<256xi8>, tensor<512xi8>, tensor<512xi8>, tensor<32768xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_001_qwen_attention_metadata_case> @qwen3_30b_a3b_q4_k_m_pp256_001_qwen_attention_metadata - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_002_ggml_get_rows_f32_case { - %token_count = check.literal value(64) : index - %row_count = check.literal value(151936) : index - %hidden_size = check.literal value(2048) : index - %token_ids = check.generate.fill value(0) : tensor<256xi8> - %weight = check.generate.fill value(0) : tensor<175030272xi8> - %output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<256xi8>, tensor<175030272xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_002_ggml_get_rows_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_002_ggml_get_rows_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_003_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(2.0) : tensor<64x2048xf32> - %rhs = check.generate.fill value(3.0) : tensor<2048xf32> - %output = check.generate.fill value(0.0) : tensor<64x2048xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<64x2048xf32>, tensor<2048xf32>, tensor<64x2048xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_003_ggml_rmsnorm_binary_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_003_ggml_rmsnorm_binary_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_004_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %output = check.generate.fill value(0) : tensor<1048576xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<4718592xi8>, tensor<1048576xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_004_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_004_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_005_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<860160xi8> - %output = check.generate.fill value(0) : tensor<131072xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<860160xi8>, tensor<131072xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_005_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_005_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_006_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<589824xi8> - %output = check.generate.fill value(0) : tensor<131072xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<589824xi8>, tensor<131072xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_006_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_006_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_007_qwen3_moe_attention_postprocess_f32_f16_case { - %token_count = check.literal value(64) : index - %cache_row_count = check.literal value(256) : index - %positions = check.generate.fill value(0) : tensor<256xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %query_input = check.generate.fill value(0) : tensor<1048576xi8> - %key_input = check.generate.fill value(0) : tensor<131072xi8> - %value_input = check.generate.fill value(0) : tensor<131072xi8> - %query_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %key_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %inverse_frequencies = check.generate.fill value(0) : tensor<256xi8> - %query_output = check.generate.fill value(0) : tensor<1048576xi8> - %key_cache = check.generate.fill value(0) : tensor<262144xi8> - %value_cache = check.generate.fill value(0) : tensor<262144xi8> - kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<256xi8>, tensor<512xi8>, tensor<512xi8>, tensor<1048576xi8>, tensor<131072xi8>, tensor<131072xi8>, tensor<512xi8>, tensor<512xi8>, tensor<256xi8>, tensor<1048576xi8>, tensor<262144xi8>, tensor<262144xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_007_qwen3_moe_attention_postprocess_f32_f16_case> @qwen3_30b_a3b_q4_k_m_pp256_007_qwen3_moe_attention_postprocess_f32_f16 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_008_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<64x32x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x256xf16> - %output = check.generate.fill value(1.0) : tensor<64x32x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<64x256xf16>, tensor<64x32x128xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_008_ggml_flash_attention_f32_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_008_ggml_flash_attention_f32_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_009_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<1048576xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %residual_input = check.generate.fill value(0) : tensor<524288xi8> - %residual_output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<1048576xi8>, tensor<4718592xi8>, tensor<524288xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_009_ggml_mul_mat_add_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_009_ggml_mul_mat_add_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_010_qwen3_moe_router_projection_f32_four_row_wave32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<1048576xi8> - %output = check.generate.fill value(0) : tensor<32768xi8> - kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<1048576xi8>, tensor<32768xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_010_qwen3_moe_router_projection_f32_four_row_wave32_case> @qwen3_30b_a3b_q4_k_m_pp256_010_qwen3_moe_router_projection_f32_four_row_wave32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_011_qwen3_moe_router_top8_f32_case { - %token_count = check.literal value(64) : index - %route_id_stride = check.literal value(128) : index - %logits = check.generate.fill value(0) : tensor<32768xi8> - %route_ids = check.generate.fill value(0) : tensor<32768xi8> - %route_weights = check.generate.fill value(0) : tensor<2048xi8> - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<32768xi8>, tensor<32768xi8>, tensor<2048xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_011_qwen3_moe_router_top8_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_011_qwen3_moe_router_top8_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_012_ggml_moe_build_expert_table_case { - %token_count = check.literal value(64) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.fill value(0) : tensor<32768xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<32768xi8>, tensor<33280xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_012_ggml_moe_build_expert_table_case> @qwen3_30b_a3b_q4_k_m_pp256_012_ggml_moe_build_expert_table - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_013_ggml_moe_build_expert_partition_table_case { - %token_count = check.literal value(64) : index - %route_count = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %partition_table = check.generate.fill value(0) : tensor<580xi8> - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<33280xi8>, tensor<580xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_013_ggml_moe_build_expert_partition_table_case> @qwen3_30b_a3b_q4_k_m_pp256_013_ggml_moe_build_expert_partition_table - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %partition_table = check.generate.fill value(0) : tensor<580xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<786432xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : [index](index, tensor<524288xi8>, tensor<33280xi8>, tensor<580xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<786432xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_015_ggml_mul_mat_id_f16_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<786432xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %weight = check.generate.fill value(0) : tensor<165150720xi8> - %output = check.generate.fill value(0) : tensor<2097152xi8> - kernel.launch @ggml_mul_mat_id_f16_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %output) : [index](index, tensor<786432xi8>, tensor<33280xi8>, tensor<165150720xi8>, tensor<2097152xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_015_ggml_mul_mat_id_f16_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_015_ggml_mul_mat_id_f16_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_case { - %token_count = check.literal value(64) : index - %route_weights = check.generate.fill value(0) : tensor<2048xi8> - %routed_output = check.generate.fill value(0) : tensor<2097152xi8> - %hidden_state = check.generate.fill value(0) : tensor<524288xi8> - %next_norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %next_projection_input = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<2048xi8>, tensor<2097152xi8>, tensor<524288xi8>, tensor<8192xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_017_ggml_mul_mat_id_f16_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<786432xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<2097152xi8> - kernel.launch @ggml_mul_mat_id_f16_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %output) : [index](index, tensor<786432xi8>, tensor<33280xi8>, tensor<113246208xi8>, tensor<2097152xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_017_ggml_mul_mat_id_f16_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_017_ggml_mul_mat_id_f16_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_018_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<1048576xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<1048576xi8>, tensor<4718592xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_018_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp256_018_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_019_ggml_gather_add_f32_case { - %source_token_count = check.literal value(64) : index - %output_token_count = check.literal value(1) : index - %hidden_size = check.literal value(2048) : index - %attention = check.generate.fill value(0.0) : tensor<64x2048xf32> - %residual = check.generate.fill value(0.0) : tensor<64x2048xf32> - %output_ids = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<64x2048xf32>, tensor<64x2048xf32>, tensor<1xi32>, tensor<1x2048xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_019_ggml_gather_add_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_019_ggml_gather_add_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<8192xi8> - %normalized_output = check.generate.fill value(0) : tensor<8192xi8> - %q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %normalized_output, %q8_output) : [index](index, tensor<8192xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_pp256_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_021_qwen3_moe_router_projection_top8_fused_decode_f32_case { - %token_count = check.literal value(1) : index - %route_id_stride = check.literal value(128) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<1048576xi8> - %logits = check.generate.fill value(0) : tensor<512xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %logits, %completion_counter, %route_ids, %route_weights) : [index, index](index, index, tensor<8192xi8>, tensor<1048576xi8>, tensor<512xi8>, tensor<4xi8>, tensor<512xi8>, tensor<32xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_021_qwen3_moe_router_projection_top8_fused_decode_f32_case> @qwen3_30b_a3b_q4_k_m_pp256_021_qwen3_moe_router_projection_top8_fused_decode_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<24576xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<512xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<24576xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case> @qwen3_30b_a3b_q4_k_m_pp256_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %input = check.generate.fill value(0) : tensor<24576xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - %weight = check.generate.fill value(0) : tensor<165150720xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - %norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<24576xi8>, tensor<512xi8>, tensor<32xi8>, tensor<165150720xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<4xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case> @qwen3_30b_a3b_q4_k_m_pp256_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_pp256_024_ggml_linear_q6k_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(151936) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %weight = check.generate.fill value(0) : tensor<255252480xi8> - %output = check.generate.fill value(0) : tensor<607744xi8> - kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<2304xi8>, tensor<255252480xi8>, tensor<607744xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp256_024_ggml_linear_q6k_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_pp256_024_ggml_linear_q6k_q8_1_x4 - -// Scenario: pp512 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_000_qwen_attention_context_base_capture_case { - %positions = check.generate.fill value(0) : tensor<256xi8> - %control = check.generate.fill value(0) : tensor<4xi8> - kernel.launch @qwen_attention_context_base_capture[](%positions, %control) : [](tensor<256xi8>, tensor<4xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_000_qwen_attention_context_base_capture_case> @qwen3_30b_a3b_q4_k_m_pp512_000_qwen_attention_context_base_capture - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_001_qwen_attention_metadata_case { - %token_count = check.literal value(64) : index - %context_capacity = check.literal value(256) : index - %control = check.generate.fill value(0) : tensor<4xi8> - %positions = check.generate.fill value(0) : tensor<256xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %attention_mask = check.generate.fill value(0) : tensor<32768xi8> - kernel.launch @qwen_attention_metadata[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<4xi8>, tensor<256xi8>, tensor<512xi8>, tensor<512xi8>, tensor<32768xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_001_qwen_attention_metadata_case> @qwen3_30b_a3b_q4_k_m_pp512_001_qwen_attention_metadata - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_002_ggml_get_rows_f32_case { - %token_count = check.literal value(64) : index - %row_count = check.literal value(151936) : index - %hidden_size = check.literal value(2048) : index - %token_ids = check.generate.fill value(0) : tensor<256xi8> - %weight = check.generate.fill value(0) : tensor<175030272xi8> - %output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<256xi8>, tensor<175030272xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_002_ggml_get_rows_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_002_ggml_get_rows_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_003_ggml_rmsnorm_binary_f32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(2.0) : tensor<64x2048xf32> - %rhs = check.generate.fill value(3.0) : tensor<2048xf32> - %output = check.generate.fill value(0.0) : tensor<64x2048xf32> - kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<64x2048xf32>, tensor<2048xf32>, tensor<64x2048xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_003_ggml_rmsnorm_binary_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_003_ggml_rmsnorm_binary_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_004_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %output = check.generate.fill value(0) : tensor<1048576xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<4718592xi8>, tensor<1048576xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_004_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_004_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_005_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<860160xi8> - %output = check.generate.fill value(0) : tensor<131072xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<860160xi8>, tensor<131072xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_005_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_005_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_006_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<589824xi8> - %output = check.generate.fill value(0) : tensor<131072xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<589824xi8>, tensor<131072xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_006_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_006_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_007_qwen3_moe_attention_postprocess_f32_f16_case { - %token_count = check.literal value(64) : index - %cache_row_count = check.literal value(512) : index - %positions = check.generate.fill value(0) : tensor<256xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %query_input = check.generate.fill value(0) : tensor<1048576xi8> - %key_input = check.generate.fill value(0) : tensor<131072xi8> - %value_input = check.generate.fill value(0) : tensor<131072xi8> - %query_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %key_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %inverse_frequencies = check.generate.fill value(0) : tensor<256xi8> - %query_output = check.generate.fill value(0) : tensor<1048576xi8> - %key_cache = check.generate.fill value(0) : tensor<524288xi8> - %value_cache = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @qwen3_moe_attention_postprocess_f32_f16[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : [index, index](index, index, tensor<256xi8>, tensor<512xi8>, tensor<512xi8>, tensor<1048576xi8>, tensor<131072xi8>, tensor<131072xi8>, tensor<512xi8>, tensor<512xi8>, tensor<256xi8>, tensor<1048576xi8>, tensor<524288xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_007_qwen3_moe_attention_postprocess_f32_f16_case> @qwen3_30b_a3b_q4_k_m_pp512_007_qwen3_moe_attention_postprocess_f32_f16 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_008_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<64x32x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x256xf16> - %output = check.generate.fill value(1.0) : tensor<64x32x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<64x256xf16>, tensor<64x32x128xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_008_ggml_flash_attention_f32_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_008_ggml_flash_attention_f32_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_009_ggml_mul_mat_add_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<1048576xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %residual_input = check.generate.fill value(0) : tensor<524288xi8> - %residual_output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_mul_mat_add_f32_f32_wmma[%token_count](%token_count, %input, %weight, %residual_input, %residual_output) : [index](index, tensor<1048576xi8>, tensor<4718592xi8>, tensor<524288xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_009_ggml_mul_mat_add_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_009_ggml_mul_mat_add_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_010_qwen3_moe_router_projection_f32_four_row_wave32_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %weight = check.generate.fill value(0) : tensor<1048576xi8> - %output = check.generate.fill value(0) : tensor<32768xi8> - kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<524288xi8>, tensor<1048576xi8>, tensor<32768xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_010_qwen3_moe_router_projection_f32_four_row_wave32_case> @qwen3_30b_a3b_q4_k_m_pp512_010_qwen3_moe_router_projection_f32_four_row_wave32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_011_qwen3_moe_router_top8_f32_case { - %token_count = check.literal value(64) : index - %route_id_stride = check.literal value(128) : index - %logits = check.generate.fill value(0) : tensor<32768xi8> - %route_ids = check.generate.fill value(0) : tensor<32768xi8> - %route_weights = check.generate.fill value(0) : tensor<2048xi8> - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<32768xi8>, tensor<32768xi8>, tensor<2048xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_011_qwen3_moe_router_top8_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_011_qwen3_moe_router_top8_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_012_ggml_moe_build_expert_table_case { - %token_count = check.literal value(64) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.fill value(0) : tensor<32768xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<32768xi8>, tensor<33280xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_012_ggml_moe_build_expert_table_case> @qwen3_30b_a3b_q4_k_m_pp512_012_ggml_moe_build_expert_table - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_013_ggml_moe_build_expert_partition_table_case { - %token_count = check.literal value(64) : index - %route_count = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %partition_table = check.generate.fill value(0) : tensor<580xi8> - kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<33280xi8>, tensor<580xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_013_ggml_moe_build_expert_partition_table_case> @qwen3_30b_a3b_q4_k_m_pp512_013_ggml_moe_build_expert_partition_table - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<524288xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %partition_table = check.generate.fill value(0) : tensor<580xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<786432xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : [index](index, tensor<524288xi8>, tensor<33280xi8>, tensor<580xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<786432xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_015_ggml_mul_mat_id_f16_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<786432xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %weight = check.generate.fill value(0) : tensor<165150720xi8> - %output = check.generate.fill value(0) : tensor<2097152xi8> - kernel.launch @ggml_mul_mat_id_f16_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %output) : [index](index, tensor<786432xi8>, tensor<33280xi8>, tensor<165150720xi8>, tensor<2097152xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_015_ggml_mul_mat_id_f16_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_015_ggml_mul_mat_id_f16_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_case { - %token_count = check.literal value(64) : index - %route_weights = check.generate.fill value(0) : tensor<2048xi8> - %routed_output = check.generate.fill value(0) : tensor<2097152xi8> - %hidden_state = check.generate.fill value(0) : tensor<524288xi8> - %next_norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %next_projection_input = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<2048xi8>, tensor<2097152xi8>, tensor<524288xi8>, tensor<8192xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_017_ggml_mul_mat_id_f16_f16_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<786432xi8> - %expert_table = check.generate.fill value(0) : tensor<33280xi8> - %weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<2097152xi8> - kernel.launch @ggml_mul_mat_id_f16_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %output) : [index](index, tensor<786432xi8>, tensor<33280xi8>, tensor<113246208xi8>, tensor<2097152xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_017_ggml_mul_mat_id_f16_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_017_ggml_mul_mat_id_f16_f16_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_018_ggml_mul_mat_f32_f32_wmma_case { - %token_count = check.literal value(64) : index - %input = check.generate.fill value(0) : tensor<1048576xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %output = check.generate.fill value(0) : tensor<524288xi8> - kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<1048576xi8>, tensor<4718592xi8>, tensor<524288xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_018_ggml_mul_mat_f32_f32_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_018_ggml_mul_mat_f32_f32_wmma - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_019_ggml_gather_add_f32_case { - %source_token_count = check.literal value(64) : index - %output_token_count = check.literal value(1) : index - %hidden_size = check.literal value(2048) : index - %attention = check.generate.fill value(0.0) : tensor<64x2048xf32> - %residual = check.generate.fill value(0.0) : tensor<64x2048xf32> - %output_ids = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<64x2048xf32>, tensor<64x2048xf32>, tensor<1xi32>, tensor<1x2048xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_019_ggml_gather_add_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_019_ggml_gather_add_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<8192xi8> - %normalized_output = check.generate.fill value(0) : tensor<8192xi8> - %q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %normalized_output, %q8_output) : [index](index, tensor<8192xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_pp512_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_021_qwen3_moe_router_projection_top8_fused_decode_f32_case { - %token_count = check.literal value(1) : index - %route_id_stride = check.literal value(128) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<1048576xi8> - %logits = check.generate.fill value(0) : tensor<512xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %logits, %completion_counter, %route_ids, %route_weights) : [index, index](index, index, tensor<8192xi8>, tensor<1048576xi8>, tensor<512xi8>, tensor<4xi8>, tensor<512xi8>, tensor<32xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_021_qwen3_moe_router_projection_top8_fused_decode_f32_case> @qwen3_30b_a3b_q4_k_m_pp512_021_qwen3_moe_router_projection_top8_fused_decode_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<24576xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<512xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<24576xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case> @qwen3_30b_a3b_q4_k_m_pp512_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %input = check.generate.fill value(0) : tensor<24576xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - %weight = check.generate.fill value(0) : tensor<165150720xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - %norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<24576xi8>, tensor<512xi8>, tensor<32xi8>, tensor<165150720xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<4xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case> @qwen3_30b_a3b_q4_k_m_pp512_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_024_ggml_linear_q6k_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(151936) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %weight = check.generate.fill value(0) : tensor<255252480xi8> - %output = check.generate.fill value(0) : tensor<607744xi8> - kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<2304xi8>, tensor<255252480xi8>, tensor<607744xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_024_ggml_linear_q6k_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_pp512_024_ggml_linear_q6k_q8_1_x4 - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_025_qwen_attention_metadata_case { - %token_count = check.literal value(64) : index - %context_capacity = check.literal value(512) : index - %control = check.generate.fill value(0) : tensor<4xi8> - %positions = check.generate.fill value(0) : tensor<256xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<512xi8> - %attention_mask = check.generate.fill value(0) : tensor<65536xi8> - kernel.launch @qwen_attention_metadata[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<4xi8>, tensor<256xi8>, tensor<512xi8>, tensor<512xi8>, tensor<65536xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_025_qwen_attention_metadata_case> @qwen3_30b_a3b_q4_k_m_pp512_025_qwen_attention_metadata - -check.case public @qwen3_30b_a3b_q4_k_m_pp512_026_ggml_flash_attention_f32_f16_wmma_case { - %query_token_count = check.literal value(64) : index - %key_value_token_count = check.literal value(512) : index - %query = check.generate.fill value(0.0) : tensor<64x32x128xf32> - %key = check.generate.fill value(0.0) : tensor<512x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<512x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<64x512xf16> - %output = check.generate.fill value(1.0) : tensor<64x32x128xf32> - kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<64x32x128xf32>, tensor<512x4x128xf16>, tensor<512x4x128xf16>, tensor<64x512xf16>, tensor<64x32x128xf32>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_pp512_026_ggml_flash_attention_f32_f16_wmma_case> @qwen3_30b_a3b_q4_k_m_pp512_026_ggml_flash_attention_f32_f16_wmma - -// Scenario: tg8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_000_qwen_attention_context_base_capture_case { - %positions = check.generate.fill value(0) : tensor<4xi8> - %control = check.generate.fill value(0) : tensor<4xi8> - kernel.launch @qwen_attention_context_base_capture[](%positions, %control) : [](tensor<4xi8>, tensor<4xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_000_qwen_attention_context_base_capture_case> @qwen3_30b_a3b_q4_k_m_tg8_000_qwen_attention_context_base_capture - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_001_qwen_attention_metadata_case { - %token_count = check.literal value(1) : index - %context_capacity = check.literal value(256) : index - %control = check.generate.fill value(0) : tensor<4xi8> - %positions = check.generate.fill value(0) : tensor<4xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %attention_mask = check.generate.fill value(0) : tensor<512xi8> - kernel.launch @qwen_attention_metadata[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<4xi8>, tensor<4xi8>, tensor<8xi8>, tensor<8xi8>, tensor<512xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_001_qwen_attention_metadata_case> @qwen3_30b_a3b_q4_k_m_tg8_001_qwen_attention_metadata - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_002_ggml_get_rows_f32_case { - %token_count = check.literal value(1) : index - %row_count = check.literal value(151936) : index - %hidden_size = check.literal value(2048) : index - %token_ids = check.generate.fill value(0) : tensor<4xi8> - %weight = check.generate.fill value(0) : tensor<175030272xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi8>, tensor<175030272xi8>, tensor<8192xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_002_ggml_get_rows_f32_case> @qwen3_30b_a3b_q4_k_m_tg8_002_ggml_get_rows_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_003_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<8192xi8> - %normalized_output = check.generate.fill value(0) : tensor<8192xi8> - %q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4[%token_count](%token_count, %input, %weight, %normalized_output, %q8_output) : [index](index, tensor<8192xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_003_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_tg8_003_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_004_qwen3_moe_attention_qkv_postprocess_fused_decode_case { - %token_count = check.literal value(1) : index - %cache_row_count = check.literal value(256) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %query_weight = check.generate.fill value(0) : tensor<4718592xi8> - %key_weight = check.generate.fill value(0) : tensor<589824xi8> - %value_weight = check.generate.fill value(0) : tensor<860160xi8> - %positions = check.generate.fill value(0) : tensor<4xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %query_output_raw = check.generate.fill value(0) : tensor<16384xi8> - %key_output_raw = check.generate.fill value(0) : tensor<2048xi8> - %value_output_raw = check.generate.fill value(0) : tensor<2048xi8> - %query_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %key_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %inverse_frequencies = check.generate.fill value(0) : tensor<256xi8> - %query_output = check.generate.fill value(0) : tensor<16384xi8> - %key_cache = check.generate.fill value(0) : tensor<262144xi8> - %value_cache = check.generate.fill value(0) : tensor<262144xi8> - %completion_counters = check.generate.fill value(0) : tensor<160xi8> - kernel.launch @qwen3_moe_attention_qkv_postprocess_fused_decode[%token_count, %cache_row_count](%token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : [index, index](index, index, tensor<2304xi8>, tensor<4718592xi8>, tensor<589824xi8>, tensor<860160xi8>, tensor<4xi8>, tensor<8xi8>, tensor<8xi8>, tensor<16384xi8>, tensor<2048xi8>, tensor<2048xi8>, tensor<512xi8>, tensor<512xi8>, tensor<256xi8>, tensor<16384xi8>, tensor<262144xi8>, tensor<262144xi8>, tensor<160xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_004_qwen3_moe_attention_qkv_postprocess_fused_decode_case> @qwen3_30b_a3b_q4_k_m_tg8_004_qwen3_moe_attention_qkv_postprocess_fused_decode - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_005_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8_case { - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<256x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<256xf16> - %partial_max = check.generate.fill value(0.0) : tensor<4x4x16xf32> - %partial_sum = check.generate.fill value(0.0) : tensor<4x4x16xf32> - %partial_output = check.generate.fill value(0.0) : tensor<4x4x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<4xi32> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %next_q8_output = check.generate.fill value(0) : tensor<4608xi8> - kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %next_q8_output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_005_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_005_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_006_qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_case { - %token_count = check.literal value(1) : index - %q8_input = check.generate.fill value(0) : tensor<4608xi8> - %weight = check.generate.fill value(0) : tensor<4718592xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - %norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %normalized_output = check.generate.fill value(0) : tensor<8192xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8[%token_count](%token_count, %q8_input, %weight, %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output) : [index](index, tensor<4608xi8>, tensor<4718592xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<4xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_006_qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_006_qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_007_qwen3_moe_router_projection_top8_fused_decode_f32_case { - %token_count = check.literal value(1) : index - %route_id_stride = check.literal value(128) : index - %input = check.generate.fill value(0) : tensor<8192xi8> - %weight = check.generate.fill value(0) : tensor<1048576xi8> - %logits = check.generate.fill value(0) : tensor<512xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %logits, %completion_counter, %route_ids, %route_weights) : [index, index](index, index, tensor<8192xi8>, tensor<1048576xi8>, tensor<512xi8>, tensor<4xi8>, tensor<512xi8>, tensor<32xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_007_qwen3_moe_router_projection_top8_fused_decode_f32_case> @qwen3_30b_a3b_q4_k_m_tg8_007_qwen3_moe_router_projection_top8_fused_decode_f32 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_008_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<24576xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<512xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<24576xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_008_qwen3_moe_routed_gate_up_swiglu_q4k_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_008_qwen3_moe_routed_gate_up_swiglu_q4k_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_009_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %input = check.generate.fill value(0) : tensor<24576xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - %weight = check.generate.fill value(0) : tensor<165150720xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - %norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<24576xi8>, tensor<512xi8>, tensor<32xi8>, tensor<165150720xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<4xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_009_qwen3_moe_routed_down_q6k_f32_wave64_next_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_009_qwen3_moe_routed_down_q6k_f32_wave64_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_010_qwen3_moe_attention_qkv_postprocess_fused_decode_case { - %token_count = check.literal value(1) : index - %cache_row_count = check.literal value(256) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %query_weight = check.generate.fill value(0) : tensor<4718592xi8> - %key_weight = check.generate.fill value(0) : tensor<589824xi8> - %value_weight = check.generate.fill value(0) : tensor<589824xi8> - %positions = check.generate.fill value(0) : tensor<4xi8> - %key_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %value_cache_indices = check.generate.fill value(0) : tensor<8xi8> - %query_output_raw = check.generate.fill value(0) : tensor<16384xi8> - %key_output_raw = check.generate.fill value(0) : tensor<2048xi8> - %value_output_raw = check.generate.fill value(0) : tensor<2048xi8> - %query_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %key_norm_weight = check.generate.fill value(0) : tensor<512xi8> - %inverse_frequencies = check.generate.fill value(0) : tensor<256xi8> - %query_output = check.generate.fill value(0) : tensor<16384xi8> - %key_cache = check.generate.fill value(0) : tensor<262144xi8> - %value_cache = check.generate.fill value(0) : tensor<262144xi8> - %completion_counters = check.generate.fill value(0) : tensor<160xi8> - kernel.launch @qwen3_moe_attention_qkv_postprocess_fused_decode[%token_count, %cache_row_count](%token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : [index, index](index, index, tensor<2304xi8>, tensor<4718592xi8>, tensor<589824xi8>, tensor<589824xi8>, tensor<4xi8>, tensor<8xi8>, tensor<8xi8>, tensor<16384xi8>, tensor<2048xi8>, tensor<2048xi8>, tensor<512xi8>, tensor<512xi8>, tensor<256xi8>, tensor<16384xi8>, tensor<262144xi8>, tensor<262144xi8>, tensor<160xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_010_qwen3_moe_attention_qkv_postprocess_fused_decode_case> @qwen3_30b_a3b_q4_k_m_tg8_010_qwen3_moe_attention_qkv_postprocess_fused_decode - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_011_qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %gate_weight = check.generate.fill value(0) : tensor<113246208xi8> - %up_weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<24576xi8> - %completion_counters = check.generate.fill value(0) : tensor<192xi8> - %next_q8_output = check.generate.fill value(0) : tensor<6912xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<512xi8>, tensor<113246208xi8>, tensor<113246208xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<6912xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_011_qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_011_qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_012_qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<6912xi8> - %route_ids = check.generate.fill value(0) : tensor<512xi8> - %route_weights = check.generate.fill value(0) : tensor<32xi8> - %weight = check.generate.fill value(0) : tensor<113246208xi8> - %output = check.generate.fill value(0) : tensor<8192xi8> - %norm_weight = check.generate.fill value(0) : tensor<8192xi8> - %completion_counter = check.generate.fill value(0) : tensor<4xi8> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<6912xi8>, tensor<512xi8>, tensor<32xi8>, tensor<113246208xi8>, tensor<8192xi8>, tensor<8192xi8>, tensor<4xi8>, tensor<2304xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_012_qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_case> @qwen3_30b_a3b_q4_k_m_tg8_012_qwen3_moe_routed_down_q4k_q8_1_x4_next_q8 - -check.case public @qwen3_30b_a3b_q4_k_m_tg8_013_ggml_linear_q6k_q8_1_x4_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(151936) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %weight = check.generate.fill value(0) : tensor<255252480xi8> - %output = check.generate.fill value(0) : tensor<607744xi8> - kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %output) : [index, index, index](index, index, index, tensor<2304xi8>, tensor<255252480xi8>, tensor<607744xi8>) - check.return -} - -check.benchmark<@qwen3_30b_a3b_q4_k_m_tg8_013_ggml_linear_q6k_q8_1_x4_case> @qwen3_30b_a3b_q4_k_m_tg8_013_ggml_linear_q6k_q8_1_x4 diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json deleted file mode 100644 index 912666842ec6..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp256.json +++ /dev/null @@ -1,1073 +0,0 @@ -{ - "command_count": 674, - "dispatch_count": 25, - "dispatches": [ - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_000_qwen_attention_context_base_capture", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/qwen", - "count": 1, - "integer_parameters": {}, - "kernel": "qwen:qwen_attention_context_base_capture", - "library_sources": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "primary_sources": [ - "attention_state_initialize.loom" - ], - "sources": [ - "attention_state_initialize.loom" - ], - "symbol": "qwen_attention_context_base_capture", - "workload_parameters": [] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_001_qwen_attention_metadata", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "context_capacity": 256, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen_attention_metadata", - "library_sources": [], - "primary_sources": [ - "../qwen/attention_metadata.loom" - ], - "sources": [ - "../qwen/attention_metadata.loom" - ], - "symbol": "qwen_attention_metadata", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_002_ggml_get_rows_f32", - "compile_parameters": { - "ggml.get_rows_f32.hidden_capacity": "2048", - "ggml.get_rows_f32.token_capacity": "64", - "ggml.get_rows_f32.weight_format": "4" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "hidden_size": 2048, - "row_count": 151936, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_get_rows_f32", - "library_sources": [ - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/rows/get_rows_f32.loom" - ], - "sources": [ - "ops/rows/get_rows_f32.loom", - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_get_rows_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "row_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_003_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "2048", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_004_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "4096", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_005_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "512", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "6", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 24, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_006_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "512", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 72, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_007_qwen3_moe_attention_postprocess_f32_f16", - "compile_parameters": { - "qwen3_moe.attention.head_size": "128", - "qwen3_moe.attention.key_value_size": "512", - "qwen3_moe.attention.query_size": "4096", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 48, - "integer_parameters": { - "cache_row_count": 256, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom" - ], - "sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_attention_postprocess_f32_f16", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "cache_row_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_008_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "4", - "ggml.flash_attention.query_head_count": "32" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "key_value_token_count": 256, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_009_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "4096", - "ggml.mul_mat_postops.output_size": "2048", - "ggml.mul_mat_postops.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 47, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_010_qwen3_moe_router_projection_f32_four_row_wave32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 47, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "sources": [ - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_011_qwen3_moe_router_top8_f32", - "compile_parameters": { - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.router.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 47, - "integer_parameters": { - "route_id_stride": 128, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_router_top8_f32", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/router_top8_f32.loom" - ], - "sources": [ - "qwen3_moe/router_top8_f32.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_router_top8_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_012_ggml_moe_build_expert_table", - "compile_parameters": { - "ggml.moe_routing.descriptor_expert_mask": "127", - "ggml.moe_routing.descriptor_partition_shift": "7", - "ggml.moe_routing.descriptor_row_count_shift": "13", - "ggml.moe_routing.expert_count": "128", - "ggml.moe_routing.partition_workgroup_size": "128", - "ggml.moe_routing.route_count": "8", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 47, - "integer_parameters": { - "expert_count": 128, - "route_count": 8, - "route_stride": 128, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_moe_build_expert_table", - "library_sources": [], - "primary_sources": [ - "ops/moe_routing_tables.loom" - ], - "sources": [ - "ops/moe_routing_tables.loom" - ], - "symbol": "ggml_moe_build_expert_table", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_013_ggml_moe_build_expert_partition_table", - "compile_parameters": { - "ggml.moe_routing.descriptor_expert_mask": "127", - "ggml.moe_routing.descriptor_partition_shift": "7", - "ggml.moe_routing.descriptor_row_count_shift": "13", - "ggml.moe_routing.expert_count": "128", - "ggml.moe_routing.partition_workgroup_size": "128", - "ggml.moe_routing.route_count": "8", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 47, - "integer_parameters": { - "expert_count": 128, - "route_count": 8, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_moe_build_expert_partition_table", - "library_sources": [], - "primary_sources": [ - "ops/moe_routing_tables.loom" - ], - "sources": [ - "ops/moe_routing_tables.loom" - ], - "symbol": "ggml_moe_build_expert_partition_table", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "compile_parameters": { - "qwen3_moe.routed_gate_up.expert_count": "128", - "qwen3_moe.routed_gate_up.input_size": "2048", - "qwen3_moe.routed_gate_up.output_size": "768", - "qwen3_moe.routed_gate_up.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 47, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_015_ggml_mul_mat_id_f16_f16_wmma", - "compile_parameters": { - "ggml.mul_mat_id_f16_f16.expert_count": "128", - "ggml.mul_mat_id_f16_f16.input_size": "768", - "ggml.mul_mat_id_f16_f16.output_size": "2048", - "ggml.mul_mat_id_f16_f16.route_count": "8", - "ggml.mul_mat_id_f16_f16.weight_format": "6", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 23, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_id_f16_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom" - ], - "sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom", - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "symbol": "ggml_mul_mat_id_f16_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.routed_down.expert_count": "128", - "qwen3_moe.routed_down.input_size": "768", - "qwen3_moe.routed_down.output_size": "2048", - "qwen3_moe.routed_down.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 47, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "library_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" - ], - "sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_017_ggml_mul_mat_id_f16_f16_wmma", - "compile_parameters": { - "ggml.mul_mat_id_f16_f16.expert_count": "128", - "ggml.mul_mat_id_f16_f16.input_size": "768", - "ggml.mul_mat_id_f16_f16.output_size": "2048", - "ggml.mul_mat_id_f16_f16.route_count": "8", - "ggml.mul_mat_id_f16_f16.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 24, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_id_f16_f16_wmma", - "library_sources": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "primary_sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom" - ], - "sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom", - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "symbol": "ggml_mul_mat_id_f16_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_018_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "4096", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "2048", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 1, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_019_ggml_gather_add_f32", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/hrx", - "count": 1, - "integer_parameters": { - "hidden_size": 2048, - "output_token_count": 1, - "source_token_count": 64 - }, - "kernel": "hrx:ggml_gather_add_f32", - "library_sources": [], - "primary_sources": [ - "gather_add_f32.loom" - ], - "sources": [ - "gather_add_f32.loom" - ], - "symbol": "ggml_gather_add_f32", - "workload_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "compile_parameters": { - "ggml.quantize_q8_1_x4.group_capacity": "16", - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "sources": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_021_qwen3_moe_router_projection_top8_fused_decode_f32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.router.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "route_id_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", - "library_sources": [ - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ], - "primary_sources": [ - "qwen3_moe/router_projection_top8_fused_f32.loom" - ], - "sources": [ - "qwen3_moe/router_projection_top8_fused_f32.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ], - "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "compile_parameters": { - "qwen3_moe.routed_gate_up.expert_count": "128", - "qwen3_moe.routed_gate_up.input_size": "2048", - "qwen3_moe.routed_gate_up.output_size": "768", - "qwen3_moe.routed_gate_up.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "expert_count": 128, - "output_size": 768, - "route_count": 8, - "route_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.routed_down.expert_count": "128", - "qwen3_moe.routed_down.input_size": "768", - "qwen3_moe.routed_down.output_size": "2048", - "qwen3_moe.routed_down.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "expert_count": 128, - "input_size": 768, - "output_size": 2048, - "route_count": 8, - "route_id_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "sources": [ - "qwen3_moe/routed_down_q6k.loom", - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ], - "symbol": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp256_024_ggml_linear_q6k_q8_1_x4", - "compile_parameters": { - "ggml.linear_q6k_q8_1_x4.output_capacity": "151936", - "ggml.linear_q6k_q8_1_x4.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "input_size": 2048, - "output_size": 151936, - "token_count": 1 - }, - "kernel": "qwen3_moe:ggml_linear_q6k_q8_1_x4", - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "ggml/linear_q6k_q8_1_x4.loom" - ], - "sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "ggml_linear_q6k_q8_1_x4", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - } - ], - "generated_count": 25, - "kernel_counts": { - "hrx:ggml_gather_add_f32": 1, - "loom_libs:ggml_flash_attention_f32_f16_wmma": 48, - "loom_libs:ggml_get_rows_f32": 1, - "loom_libs:ggml_moe_build_expert_partition_table": 47, - "loom_libs:ggml_moe_build_expert_table": 47, - "loom_libs:ggml_mul_mat_add_f32_f32_wmma": 47, - "loom_libs:ggml_mul_mat_f32_f32_wmma": 145, - "loom_libs:ggml_mul_mat_id_f16_f16_wmma": 47, - "loom_libs:ggml_rmsnorm_binary_f32": 48, - "qwen3_moe:ggml_linear_q6k_q8_1_x4": 1, - "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16": 48, - "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4": 1, - "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8": 1, - "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32": 47, - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma": 47, - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8": 1, - "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32": 47, - "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32": 1, - "qwen3_moe:qwen3_moe_router_top8_f32": 47, - "qwen3_moe:qwen_attention_metadata": 1, - "qwen:qwen_attention_context_base_capture": 1 - }, - "loom_source": "benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom", - "model": "qwen3_30b_a3b_q4_k_m", - "scenario": "pp256", - "schema": "ggml-hrx-model-loom-benchmarks-v2" -} diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json deleted file mode 100644 index d5eaa64c859b..000000000000 --- a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.pp512.json +++ /dev/null @@ -1,1135 +0,0 @@ -{ - "command_count": 1348, - "dispatch_count": 27, - "dispatches": [ - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_000_qwen_attention_context_base_capture", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/qwen", - "count": 2, - "integer_parameters": {}, - "kernel": "qwen:qwen_attention_context_base_capture", - "library_sources": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "primary_sources": [ - "attention_state_initialize.loom" - ], - "sources": [ - "attention_state_initialize.loom" - ], - "symbol": "qwen_attention_context_base_capture", - "workload_parameters": [] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_001_qwen_attention_metadata", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "context_capacity": 256, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen_attention_metadata", - "library_sources": [], - "primary_sources": [ - "../qwen/attention_metadata.loom" - ], - "sources": [ - "../qwen/attention_metadata.loom" - ], - "symbol": "qwen_attention_metadata", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_002_ggml_get_rows_f32", - "compile_parameters": { - "ggml.get_rows_f32.hidden_capacity": "2048", - "ggml.get_rows_f32.token_capacity": "64", - "ggml.get_rows_f32.weight_format": "4" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "hidden_size": 2048, - "row_count": 151936, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_get_rows_f32", - "library_sources": [ - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/rows/get_rows_f32.loom" - ], - "sources": [ - "ops/rows/get_rows_f32.loom", - "motifs/dequant.loom", - "motifs/publish_f32.loom", - "motifs/quantize_q8_1_x4.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_get_rows_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "row_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_003_ggml_rmsnorm_binary_f32", - "compile_parameters": { - "ggml.rmsnorm_binary_f32.hidden_size": "2048", - "ggml.rmsnorm_binary_f32.op": "2", - "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 96, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_rmsnorm_binary_f32", - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "primary_sources": [ - "ops/normalization/rmsnorm.loom" - ], - "sources": [ - "ops/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm.loom", - "motifs/binary_f32_apply.loom" - ], - "symbol": "ggml_rmsnorm_binary_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_004_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "4096", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 96, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_005_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "512", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "6", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_006_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "2048", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "512", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 144, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_007_qwen3_moe_attention_postprocess_f32_f16", - "compile_parameters": { - "qwen3_moe.attention.head_size": "128", - "qwen3_moe.attention.key_value_size": "512", - "qwen3_moe.attention.query_size": "4096", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 96, - "integer_parameters": { - "cache_row_count": 512, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom" - ], - "sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_attention_postprocess_f32_f16", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "cache_row_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_008_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "4", - "ggml.flash_attention.query_head_count": "32" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "key_value_token_count": 256, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_009_ggml_mul_mat_add_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat_postops.input_size": "4096", - "ggml.mul_mat_postops.output_size": "2048", - "ggml.mul_mat_postops.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 94, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_010_qwen3_moe_router_projection_f32_four_row_wave32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 94, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "sources": [ - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_011_qwen3_moe_router_top8_f32", - "compile_parameters": { - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.router.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 94, - "integer_parameters": { - "route_id_stride": 128, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_router_top8_f32", - "library_sources": [ - "qwen3_moe/model_config.loom" - ], - "primary_sources": [ - "qwen3_moe/router_top8_f32.loom" - ], - "sources": [ - "qwen3_moe/router_top8_f32.loom", - "qwen3_moe/model_config.loom" - ], - "symbol": "qwen3_moe_router_top8_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_012_ggml_moe_build_expert_table", - "compile_parameters": { - "ggml.moe_routing.descriptor_expert_mask": "127", - "ggml.moe_routing.descriptor_partition_shift": "7", - "ggml.moe_routing.descriptor_row_count_shift": "13", - "ggml.moe_routing.expert_count": "128", - "ggml.moe_routing.partition_workgroup_size": "128", - "ggml.moe_routing.route_count": "8", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 94, - "integer_parameters": { - "expert_count": 128, - "route_count": 8, - "route_stride": 128, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_moe_build_expert_table", - "library_sources": [], - "primary_sources": [ - "ops/moe_routing_tables.loom" - ], - "sources": [ - "ops/moe_routing_tables.loom" - ], - "symbol": "ggml_moe_build_expert_table", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_013_ggml_moe_build_expert_partition_table", - "compile_parameters": { - "ggml.moe_routing.descriptor_expert_mask": "127", - "ggml.moe_routing.descriptor_partition_shift": "7", - "ggml.moe_routing.descriptor_row_count_shift": "13", - "ggml.moe_routing.expert_count": "128", - "ggml.moe_routing.partition_workgroup_size": "128", - "ggml.moe_routing.route_count": "8", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 94, - "integer_parameters": { - "expert_count": 128, - "route_count": 8, - "token_count": 64 - }, - "kernel": "loom_libs:ggml_moe_build_expert_partition_table", - "library_sources": [], - "primary_sources": [ - "ops/moe_routing_tables.loom" - ], - "sources": [ - "ops/moe_routing_tables.loom" - ], - "symbol": "ggml_moe_build_expert_partition_table", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_014_qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "compile_parameters": { - "qwen3_moe.routed_gate_up.expert_count": "128", - "qwen3_moe.routed_gate_up.input_size": "2048", - "qwen3_moe.routed_gate_up.output_size": "768", - "qwen3_moe.routed_gate_up.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 94, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_015_ggml_mul_mat_id_f16_f16_wmma", - "compile_parameters": { - "ggml.mul_mat_id_f16_f16.expert_count": "128", - "ggml.mul_mat_id_f16_f16.input_size": "768", - "ggml.mul_mat_id_f16_f16.output_size": "2048", - "ggml.mul_mat_id_f16_f16.route_count": "8", - "ggml.mul_mat_id_f16_f16.weight_format": "6", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 46, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_id_f16_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom" - ], - "sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom", - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "symbol": "ggml_mul_mat_id_f16_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_016_qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.routed_down.expert_count": "128", - "qwen3_moe.routed_down.input_size": "768", - "qwen3_moe.routed_down.output_size": "2048", - "qwen3_moe.routed_down.route_count": "8", - "qwen3_moe.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 94, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "library_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" - ], - "sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_017_ggml_mul_mat_id_f16_f16_wmma", - "compile_parameters": { - "ggml.mul_mat_id_f16_f16.expert_count": "128", - "ggml.mul_mat_id_f16_f16.input_size": "768", - "ggml.mul_mat_id_f16_f16.output_size": "2048", - "ggml.mul_mat_id_f16_f16.route_count": "8", - "ggml.mul_mat_id_f16_f16.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_id_f16_f16_wmma", - "library_sources": [ - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "primary_sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom" - ], - "sources": [ - "ops/mul_mat_id_f16_f16_wmma.loom", - "motifs/mul_mat_id_f16_f16_wmma_core.loom" - ], - "symbol": "ggml_mul_mat_id_f16_f16_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_018_ggml_mul_mat_f32_f32_wmma", - "compile_parameters": { - "ggml.mul_mat.input_size": "4096", - "ggml.mul_mat.output_accumulation": "0", - "ggml.mul_mat.output_size": "2048", - "ggml.mul_mat.output_unary_op": "23", - "ggml.mul_mat.weight_format": "4", - "ggml.workload.token_capacity": "64" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 2, - "integer_parameters": { - "token_count": 64 - }, - "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "sources": [ - "ops/mul_mat_f32_f32_wmma.loom", - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/f32_f16.loom" - ], - "symbol": "ggml_mul_mat_f32_f32_wmma", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_019_ggml_gather_add_f32", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/hrx", - "count": 2, - "integer_parameters": { - "hidden_size": 2048, - "output_token_count": 1, - "source_token_count": 64 - }, - "kernel": "hrx:ggml_gather_add_f32", - "library_sources": [], - "primary_sources": [ - "gather_add_f32.loom" - ], - "sources": [ - "gather_add_f32.loom" - ], - "symbol": "ggml_gather_add_f32", - "workload_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_020_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "compile_parameters": { - "ggml.quantize_q8_1_x4.group_capacity": "16", - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 2, - "integer_parameters": { - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "sources": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_021_qwen3_moe_router_projection_top8_fused_decode_f32", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.router.expert_count": "128", - "qwen3_moe.router.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 2, - "integer_parameters": { - "route_id_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", - "library_sources": [ - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ], - "primary_sources": [ - "qwen3_moe/router_projection_top8_fused_f32.loom" - ], - "sources": [ - "qwen3_moe/router_projection_top8_fused_f32.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ], - "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_022_qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "compile_parameters": { - "qwen3_moe.routed_gate_up.expert_count": "128", - "qwen3_moe.routed_gate_up.input_size": "2048", - "qwen3_moe.routed_gate_up.output_size": "768", - "qwen3_moe.routed_gate_up.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 2, - "integer_parameters": { - "expert_count": 128, - "output_size": 768, - "route_count": 8, - "route_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_023_qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "compile_parameters": { - "qwen3_moe.model.hidden_size": "2048", - "qwen3_moe.model.rms_epsilon": "0.000001", - "qwen3_moe.routed_down.expert_count": "128", - "qwen3_moe.routed_down.input_size": "768", - "qwen3_moe.routed_down.output_size": "2048", - "qwen3_moe.routed_down.route_count": "8", - "qwen3_moe.workload.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 2, - "integer_parameters": { - "expert_count": 128, - "input_size": 768, - "output_size": 2048, - "route_count": 8, - "route_id_stride": 128, - "token_count": 1 - }, - "kernel": "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ], - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "sources": [ - "qwen3_moe/routed_down_q6k.loom", - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ], - "symbol": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_024_ggml_linear_q6k_q8_1_x4", - "compile_parameters": { - "ggml.linear_q6k_q8_1_x4.output_capacity": "151936", - "ggml.linear_q6k_q8_1_x4.token_capacity": "1" - }, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 2, - "integer_parameters": { - "input_size": 2048, - "output_size": 151936, - "token_count": 1 - }, - "kernel": "qwen3_moe:ggml_linear_q6k_q8_1_x4", - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "primary_sources": [ - "ggml/linear_q6k_q8_1_x4.loom" - ], - "sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ], - "symbol": "ggml_linear_q6k_q8_1_x4", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_025_qwen_attention_metadata", - "compile_parameters": {}, - "corpus_dir": "kernel-corpus/kernels/qwen_moe", - "count": 1, - "integer_parameters": { - "context_capacity": 512, - "token_count": 64 - }, - "kernel": "qwen3_moe:qwen_attention_metadata", - "library_sources": [], - "primary_sources": [ - "../qwen/attention_metadata.loom" - ], - "sources": [ - "../qwen/attention_metadata.loom" - ], - "symbol": "qwen_attention_metadata", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ] - }, - { - "benchmark": "@qwen3_30b_a3b_q4_k_m_pp512_026_ggml_flash_attention_f32_f16_wmma", - "compile_parameters": { - "ggml.flash_attention.head_size": "128", - "ggml.flash_attention.key_value_head_count": "4", - "ggml.flash_attention.query_head_count": "32" - }, - "corpus_dir": "kernel-corpus/kernels/loom-libs", - "count": 48, - "integer_parameters": { - "key_value_token_count": 512, - "query_token_count": 64 - }, - "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", - "library_sources": [], - "primary_sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "sources": [ - "ops/flash_attention_f32_f16_wmma.loom" - ], - "symbol": "ggml_flash_attention_f32_f16_wmma", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ] - } - ], - "generated_count": 27, - "kernel_counts": { - "hrx:ggml_gather_add_f32": 2, - "loom_libs:ggml_flash_attention_f32_f16_wmma": 96, - "loom_libs:ggml_get_rows_f32": 2, - "loom_libs:ggml_moe_build_expert_partition_table": 94, - "loom_libs:ggml_moe_build_expert_table": 94, - "loom_libs:ggml_mul_mat_add_f32_f32_wmma": 94, - "loom_libs:ggml_mul_mat_f32_f32_wmma": 290, - "loom_libs:ggml_mul_mat_id_f16_f16_wmma": 94, - "loom_libs:ggml_rmsnorm_binary_f32": 96, - "qwen3_moe:ggml_linear_q6k_q8_1_x4": 2, - "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16": 96, - "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4": 2, - "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8": 2, - "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32": 94, - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma": 94, - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8": 2, - "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32": 94, - "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32": 2, - "qwen3_moe:qwen3_moe_router_top8_f32": 94, - "qwen3_moe:qwen_attention_metadata": 2, - "qwen:qwen_attention_context_base_capture": 2 - }, - "loom_source": "benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom", - "model": "qwen3_30b_a3b_q4_k_m", - "scenario": "pp512", - "schema": "ggml-hrx-model-loom-benchmarks-v2" -} diff --git a/ggml/src/ggml-hrx/dispatch/command-program.cpp b/ggml/src/ggml-hrx/dispatch/command-program.cpp index f789c5b09e62..458930a2e8a9 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program.cpp @@ -4,6 +4,7 @@ #include "transient-allocator.h" #include +#include #include #include #include @@ -22,6 +23,11 @@ static std::string string_value(const char * value) { return value != nullptr ? value : ""; } +static bool replay_seeded_transients_enabled() { + const char * value = std::getenv("GGML_HRX_REPLAY_SEEDED_TRANSIENTS"); + return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0'); +} + static CommandBindingOrigin command_binding_origin(const Graph & graph, const Value & value) { if (value.kind == ValueKind::External) { return CommandBindingOrigin::GraphValue; @@ -202,7 +208,7 @@ static void verify_command_list(const std::vector & commands, for (size_t i = 0; i < commands.size(); ++i) { const Command & command = commands[i]; const std::string command_context = format_command(command); - if (command.ordinal != i) { + if (command.ordinal != i && !replay_seeded_transients_enabled()) { status.log("%s has non-contiguous ordinal at index %zu", command_context.c_str(), i); } if (command.kind != CommandKind::Kernel) { @@ -400,6 +406,10 @@ static int32_t find_later_transient_writer(const std::vector & commands } static void verify_transient_reads_are_defined(const CommandProgram & program, Status & status) { + if (replay_seeded_transients_enabled()) { + return; + } + std::unordered_set defined; for (const Command & command : program.initialization_commands) { for (const CommandBinding & binding : command.bindings) { 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 5c55a3b4a357..b7e16b6e2b90 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 @@ -5,6 +5,7 @@ #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" +#include #include #include #include @@ -19,18 +20,18 @@ static constexpr KernelCatalogRef kMulMatTiledF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32"); static constexpr KernelCatalogRef kMulMatTiledF32F16AlternateKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate"); -static constexpr KernelCatalogRef kMulMatF32F32NarrowSplitK4Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_narrow_split_k4"); +static constexpr KernelCatalogRef kMulMatSkinnyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatSkinnyBiasF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_bias_publish_f32"); +static constexpr KernelCatalogRef kMulMatSkinnyAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatSkinnyBiasAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"); static constexpr KernelCatalogRef kMulMatF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_decode_wave64"); static constexpr KernelCatalogRef kMulMatAddF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_decode_wave64"); -static constexpr KernelCatalogRef kMulMatBiasF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_bias_f32_f32_wmma"); -static constexpr KernelCatalogRef kMulMatAddF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_wmma"); -static constexpr KernelCatalogRef kMulMatBiasAddF32F32WmmaKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_bias_add_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatVectorF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_f32_f32"); static constexpr KernelCatalogRef kMulMatVectorBiasF32F32Kernel = @@ -94,6 +95,10 @@ static constexpr KernelCatalogRef kMulMatVectorQ6SelectedOutputsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q6_selected_outputs_f32_f32"); static constexpr int64_t kQwenVectorVocabularyCount = 151936; +static bool deprecated_mul_mat_dispatch_disabled() { + return std::getenv("GGML_HRX_DISABLE_DEPRECATED_MUL_MAT_DISPATCH") != nullptr; +} + struct MulMatPostOpsMatch { const Value * input = nullptr; const Value * weight = nullptr; @@ -145,27 +150,33 @@ struct VectorPublishPlan { size_t q8_byte_count = 0; }; -static bool common_mul_mat_uses_narrow_split_k_route(const CommonMulMatMatch & match) { - return match.weight_format == CommonMulMatWeightFormat::Q4K && !match.has_fused_unary && - match.input_size >= 4096 && match.input_size % 1024 == 0 && match.output_size <= 64 && - match.output_size % 4 == 0 && match.token_count >= 128 && match.token_count % 32 == 0 && - common_ceil_div(match.token_count, 32) < 32; +static bool common_mul_mat_is_low_token_batch(int64_t token_count) { + return token_count >= 2 && token_count <= 5; } -static bool common_mul_mat_is_direct_f32_tiled_route(const CommonMulMatMatch & match) { - return match.matched() && match.token_count > 5 && !common_mul_mat_uses_narrow_split_k_route(match); +static bool common_mul_mat_has_skinny_reduction_shape(int64_t input_size) { + return input_size % 256 == 0; } static bool common_mul_mat_is_skinny_route(const CommonMulMatMatch & match) { - return match.matched() && match.token_count <= 5; + return match.matched() && common_mul_mat_is_low_token_batch(match.token_count) && + common_mul_mat_has_skinny_reduction_shape(match.input_size); } -static bool common_mul_mat_postops_is_direct_f32_tiled_route(const MulMatPostOpsMatch & match) { - return match.matched() && match.layout_node == nullptr && !match.requires_q8_activation && match.token_count > 5; +// Skinny owns the aligned 2..5 token optimization. Tiled owns normal prefill +// and the low-token tail fallback shapes. +static bool common_mul_mat_is_tiled_route(const CommonMulMatMatch & match) { + return match.matched() && match.token_count > 1 && !common_mul_mat_is_skinny_route(match); } static bool common_mul_mat_postops_is_skinny_route(const MulMatPostOpsMatch & match) { - return match.matched() && match.token_count <= 5; + return match.matched() && common_mul_mat_is_low_token_batch(match.token_count) && + common_mul_mat_has_skinny_reduction_shape(match.input_size); +} + +static bool common_mul_mat_postops_is_tiled_route(const MulMatPostOpsMatch & match) { + return match.matched() && match.layout_node == nullptr && !match.requires_q8_activation && match.token_count > 1 && + !common_mul_mat_postops_is_skinny_route(match); } static bool common_mul_mat_uses_factored_tiled_kernel(KernelCatalogRef kernel) { @@ -1101,13 +1112,6 @@ static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & con return match; } -static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context) { - return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, - kMulMatBiasF32F32WmmaKernel, - kMulMatAddF32F32WmmaKernel, - kMulMatBiasAddF32F32WmmaKernel); -} - static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, bool allow_batch1_bias_residual) { return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, @@ -1160,16 +1164,13 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, const CommonMulMatMatch & match, DispatchMatch & dispatch_match, CommonQ8ActivationPolicy q8_policy = CommonQ8ActivationPolicy::ExistingAlternateOnly) { - const bool split_k = common_mul_mat_uses_narrow_split_k_route(match); const bool publish_f16_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && common_mul_mat_has_f16_alternate_consumer(context.graph, match); const bool publish_q8_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && common_mul_mat_has_q8_alternate_consumer(context.graph, match); Dispatch dispatch; - dispatch.kernel = make_kernel_specialization( - split_k ? kMulMatF32F32NarrowSplitK4Kernel : - publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : - match.kernel); + dispatch.kernel = make_kernel_specialization(publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : + match.kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", common_to_config_value(match.token_count)); @@ -1178,8 +1179,9 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, 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 && match.output_size % 64 == 0 && - match.weight->alias_source.value < 0; + 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 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; @@ -1218,21 +1220,7 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context, } else { dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); } - if (split_k) { - const ValueId partial = context.next_plan_value; - const ValueId counters(partial.value + 1); - const size_t partial_bytes = 4 * match.output->byte_count; - const uint32_t counter_count = static_cast(match.token_count / 32); - dispatch_match.transients.push_back( - { partial, "common.mul_mat.narrow.split_k_partial", partial_bytes, 256 }); - dispatch_match.completion_counter_requests.push_back( - { counters, "common.mul_mat.narrow.split_k_counters", counter_count }); - dispatch.bindings.push_back({ partial, 0, partial_bytes }); - dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - dispatch.bindings.push_back({ counters, 0, counter_count * sizeof(int32_t) }); - } else { - dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_f16_alternate) { const size_t bytes = match.output->byte_count / 2; const ValueId f16_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); @@ -1755,8 +1743,9 @@ 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 && match.output_size % 64 == 0 && - match.weight->alias_source.value < 0; + 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 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; @@ -1820,23 +1809,12 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, return true; } -static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, - kMulMatBiasF32F32WmmaKernel, - kMulMatAddF32F32WmmaKernel, - kMulMatBiasAddF32F32WmmaKernel); - if (common_mul_mat_postops_is_direct_f32_tiled_route(match) || common_mul_mat_postops_is_skinny_route(match)) { - return false; - } - return build_mul_mat_postops_dispatch(context, match, dispatch_match); -} - static bool match_skinny_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, - kMulMatBiasF32F32WmmaKernel, - kMulMatAddF32F32WmmaKernel, - kMulMatBiasAddF32F32WmmaKernel); + const MulMatPostOpsMatch match = match_mul_mat_postops(context, kMulMatSkinnyF32F32Kernel, + kMulMatSkinnyBiasF32F32Kernel, + kMulMatSkinnyAddF32F32Kernel, + kMulMatSkinnyBiasAddF32F32Kernel); if (!common_mul_mat_postops_is_skinny_route(match)) { return false; } @@ -1849,7 +1827,7 @@ static bool match_tiled_mul_mat_postops_dispatch(const DispatchMatchContext & co kMulMatTiledBiasF32F32Kernel, kMulMatTiledAddF32F32Kernel, kMulMatTiledBiasAddF32F32Kernel); - if (!common_mul_mat_postops_is_direct_f32_tiled_route(match)) { + if (!common_mul_mat_postops_is_tiled_route(match)) { return false; } return build_mul_mat_postops_dispatch(context, match, dispatch_match); @@ -1950,22 +1928,9 @@ static void build_decode_mul_mat_dispatch(const CommonMulMatMatch & match, dispatch_match.dispatches.push_back(std::move(dispatch)); } -static bool match_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); - if (!match.matched()) { - return false; - } - try_match_fused_unary(context, match); - if (common_mul_mat_is_direct_f32_tiled_route(match) || common_mul_mat_is_skinny_route(match)) { - return false; - } - return build_mul_mat_dispatch(context, match, dispatch_match); -} - static bool match_skinny_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSkinnyF32F32Kernel, false); if (!match.matched()) { return false; } @@ -1983,19 +1948,7 @@ static bool match_tiled_mul_mat_dispatch(const DispatchMatchContext & context, D return false; } try_match_fused_unary(context, match); - if (!common_mul_mat_is_direct_f32_tiled_route(match)) { - return false; - } - return build_mul_mat_dispatch(context, match, dispatch_match); -} - -static bool match_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); - if (!match.matched() || !try_match_fused_unary(context, match)) { - return false; - } - if (common_mul_mat_is_direct_f32_tiled_route(match) || common_mul_mat_is_skinny_route(match)) { + if (!common_mul_mat_is_tiled_route(match)) { return false; } return build_mul_mat_dispatch(context, match, dispatch_match); @@ -2004,7 +1957,7 @@ static bool match_mul_mat_unary_dispatch(const DispatchMatchContext & context, D static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { CommonMulMatMatch match = - common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSkinnyF32F32Kernel, false); if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_skinny_route(match)) { return false; } @@ -2016,13 +1969,16 @@ static bool match_tiled_mul_mat_unary_dispatch(const DispatchMatchContext & cont CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); if (!match.matched() || !try_match_fused_unary(context, match) || - !common_mul_mat_is_direct_f32_tiled_route(match)) { + !common_mul_mat_is_tiled_route(match)) { return false; } return build_mul_mat_dispatch(context, match, dispatch_match); } static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + if (deprecated_mul_mat_dispatch_disabled()) { + return false; + } CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32DecodeWave64Kernel, true); if (!match.matched()) { @@ -2043,6 +1999,9 @@ static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, } static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + if (deprecated_mul_mat_dispatch_disabled()) { + return false; + } const DecodeMulMatAddMatch match = match_decode_mul_mat_add(context); if (!match.matched()) { return false; @@ -2170,14 +2129,6 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_skinny_mul_mat_unary_dispatch, }); - registry.add({ - "common.mul_mat_unary.legacy_f32_f32", - GGML_OP_MUL_MAT, - DispatchMatchKind::Fused, - 290, - DispatchSource::Common, - match_mul_mat_unary_dispatch, - }); registry.add({ "common.mul_mat.vector", GGML_OP_MUL_MAT, @@ -2202,14 +2153,6 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_skinny_mul_mat_postops_dispatch, }); - registry.add({ - "common.mul_mat_postops.legacy_f32_f32", - GGML_OP_MUL_MAT, - DispatchMatchKind::Fused, - 180, - DispatchSource::Common, - match_mul_mat_postops_dispatch, - }); registry.add({ "common.mul_mat.tiled_f32_f32", GGML_OP_MUL_MAT, @@ -2226,14 +2169,6 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_skinny_mul_mat_dispatch, }); - registry.add({ - "common.mul_mat.legacy_f32_f32", - GGML_OP_MUL_MAT, - DispatchMatchKind::Fused, - 80, - DispatchSource::Common, - match_mul_mat_dispatch, - }); registry.add({ "common.mul_mat_add.skinny_f32_f32_decode", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json index 92420be1208c..5dd418409b08 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -119,21 +119,6 @@ { "path": "ops/moe_routing_tables.loom" }, - { - "path": "ops/mul_mat_bias_f32_f32_wmma.loom" - }, - { - "path": "ops/mul_mat_add_f32_f32_wmma.loom" - }, - { - "path": "ops/mul_mat_bias_add_f32_f32_wmma.loom" - }, - { - "path": "ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom" - }, - { - "path": "ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom" - }, { "path": "ops/mul_mat_swiglu_f32_f32_wmma.loom" }, @@ -164,6 +149,9 @@ { "path": "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom" }, + { + "path": "ops/matmul/skinny_input_f32_publish_f32.loom" + }, { "path": "ops/mul_mat_dual_q4_f32_decode.loom" }, @@ -284,6 +272,15 @@ { "path": "motifs/f32_f16.loom" }, + { + "path": "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul/skinny_postprocess_f32.loom" + }, + { + "path": "motifs/matmul/skinny_publish_f32.loom" + }, { "path": "ops/matmul/tiled_input_f32_publish_f32.loom" }, @@ -1902,73 +1899,6 @@ "motifs/f32_f16.loom" ] }, - { - "name": "ggml_mul_mat_f32_f32_narrow_split_k4", - "family": "loom_libs", - "symbol": "ggml_mul_mat_f32_f32_narrow_split_k4", - "source": "ops/mul_mat_f32_f32_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "partial", - "output", - "completion_counters" - ], - "binding_access": [ - "read", - "read", - "read_write", - "write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_f32_f32_wmma.loom" - ], - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" - ] - }, - "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom" - ] - }, { "name": "llm_attention_q_matmul_rope_tiled_f32_f32", "family": "loom_libs", @@ -3327,49 +3257,34 @@ "compile_dependencies": [] }, { - "name": "ggml_mul_mat_add_next_rmsnorm_f32_f32_wmma", + "name": "ggml_mul_mat_dual_q4_f32_decode", "family": "loom_libs", - "symbol": "ggml_mul_mat_add_next_rmsnorm_f32_f32_wmma", - "source": "ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], + "symbol": "ggml_mul_mat_dual_q4_f32_decode", + "source": "ops/mul_mat_dual_q4_f32_decode.loom", + "workload_parameters": [], + "launch_parameters": [], "bindings": [ "input", - "weight", - "residual_input", - "residual_output", - "norm_weight", - "normalized_output", - "completion_counters" + "first_weight", + "second_weight", + "first_output", + "second_output" ], "binding_access": [ "read", "read", "read", - "read_write", - "read", - "read_write", - "read_write" + "write", + "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom" + "ops/mul_mat_dual_q4_f32_decode.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", + "ops/mul_mat_f32_f32_decode.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3381,8 +3296,7 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", + "ops/mul_mat_f32_f32_decode.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3394,30 +3308,44 @@ ] }, { - "name": "ggml_mul_mat_bias_f32_f32_wmma", + "name": "ggml_mul_mat_f32_f32_decode_wave64", "family": "loom_libs", - "symbol": "ggml_mul_mat_bias_f32_f32_wmma", - "source": "ops/mul_mat_bias_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "source": "ops/mul_mat_f32_f32_decode.loom", "workload_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" } ], "launch_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" } ], "bindings": [ "input", "weight", - "bias", "output" ], "binding_access": [ - "read", "read", "read", "read_write" @@ -3426,11 +3354,9 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_bias_f32_f32_wmma.loom" + "ops/mul_mat_f32_f32_decode.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3442,8 +3368,6 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3455,20 +3379,36 @@ ] }, { - "name": "ggml_mul_mat_add_f32_f32_wmma", + "name": "ggml_mul_mat_add_f32_f32_decode_wave64", "family": "loom_libs", - "symbol": "ggml_mul_mat_add_f32_f32_wmma", - "source": "ops/mul_mat_add_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_add_f32_f32_decode_wave64", + "source": "ops/mul_mat_add_f32_f32_decode.loom", "workload_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" } ], "launch_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" } ], "bindings": [ @@ -3487,11 +3427,10 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_add_f32_f32_wmma.loom" + "ops/mul_mat_add_f32_f32_decode.loom" ], "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", + "ops/mul_mat_f32_f32_decode.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3503,8 +3442,7 @@ ] }, "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", + "ops/mul_mat_f32_f32_decode.loom", "motifs/dequant.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", @@ -3516,10 +3454,10 @@ ] }, { - "name": "ggml_mul_mat_bias_add_f32_f32_wmma", + "name": "ggml_mul_mat_swiglu_f32_f32_wmma", "family": "loom_libs", - "symbol": "ggml_mul_mat_bias_add_f32_f32_wmma", - "source": "ops/mul_mat_bias_add_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_swiglu_f32_f32_wmma", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", "workload_parameters": [ { "name": "token_count", @@ -3534,55 +3472,59 @@ ], "bindings": [ "input", - "weight", - "bias", - "residual_input", - "residual_output" + "gate_weight", + "up_weight", + "output" ], "binding_access": [ "read", "read", "read", - "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_bias_add_f32_f32_wmma.loom" + "ops/mul_mat_swiglu_f32_f32_wmma.loom" ], "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" + "motifs/f32_f16.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/binary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" + "motifs/f32_f16.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/binary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, { - "name": "ggml_mul_mat_bias_add_next_rmsnorm_f32_f32_wmma", + "name": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", "family": "loom_libs", - "symbol": "ggml_mul_mat_bias_add_next_rmsnorm_f32_f32_wmma", - "source": "ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom", + "symbol": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", "workload_parameters": [ { "name": "token_count", @@ -3598,416 +3540,83 @@ "bindings": [ "input", "weight", - "bias", - "residual_input", - "residual_output", - "norm_weight", - "normalized_output", - "completion_counters" + "up_weight", + "output", + "f16_output" ], "binding_access": [ "read", "read", "read", - "read", - "read_write", - "read", - "read_write", - "read_write" + "write", + "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom" + "ops/mul_mat_swiglu_f32_f32_wmma.loom" ], "library_sources": [ "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" + "motifs/f32_f16.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/binary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, "compile_dependencies": [ "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/mul_mat_f32_f32_postops.loom", "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", "motifs/q4_k_f16.loom", "motifs/q6_k_f16.loom", "motifs/q8_0_f16.loom", "motifs/q8_1_f16.loom", "motifs/f16_f16.loom", "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" + "motifs/f32_f16.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/binary_f32_apply.loom", + "motifs/quantize_q8_1_x4.loom" ] }, { - "name": "ggml_mul_mat_dual_q4_f32_decode", + "name": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", "family": "loom_libs", - "symbol": "ggml_mul_mat_dual_q4_f32_decode", - "source": "ops/mul_mat_dual_q4_f32_decode.loom", - "workload_parameters": [], - "launch_parameters": [], + "symbol": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], "bindings": [ "input", - "first_weight", - "second_weight", - "first_output", - "second_output" + "gate_weight", + "up_weight", + "output" ], "binding_access": [ "read", "read", "read", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_dual_q4_f32_decode.loom" - ], - "library_sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_f32_f32_decode_wave64", - "family": "loom_libs", - "symbol": "ggml_mul_mat_f32_f32_decode_wave64", - "source": "ops/mul_mat_f32_f32_decode.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_f32_f32_decode.loom" - ], - "library_sources": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_add_f32_f32_decode_wave64", - "family": "loom_libs", - "symbol": "ggml_mul_mat_add_f32_f32_decode_wave64", - "source": "ops/mul_mat_add_f32_f32_decode.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "residual_input", - "residual_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_add_f32_f32_decode.loom" - ], - "library_sources": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - "compile_dependencies": [ - "ops/mul_mat_f32_f32_decode.loom", - "motifs/dequant.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom" - ] - }, - { - "name": "ggml_mul_mat_swiglu_f32_f32_wmma", - "family": "loom_libs", - "symbol": "ggml_mul_mat_swiglu_f32_f32_wmma", - "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom" - ], - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - { - "name": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", - "family": "loom_libs", - "symbol": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", - "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "up_weight", - "output", - "f16_output" - ], - "binding_access": [ - "read", - "read", - "read", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/mul_mat_swiglu_f32_f32_wmma.loom" - ], - "library_sources": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "motifs/mul_mat_f32_f32_wmma_core.loom", - "motifs/dequant.loom", - "motifs/unary_f32_apply.loom", - "motifs/q4_k_f16.loom", - "motifs/q6_k_f16.loom", - "motifs/q8_0_f16.loom", - "motifs/q8_1_f16.loom", - "motifs/f16_f16.loom", - "motifs/bf16_f16.loom", - "motifs/f32_f16.loom", - "motifs/mul_mat_quantized_f16_prefill.loom", - "motifs/binary_f32_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - { - "name": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - "family": "loom_libs", - "symbol": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" + "read_write" ], "target_selector": "", "compile_recipe": { @@ -6561,6 +6170,266 @@ "motifs/quantize_q8_1_x4.loom" ] }, + { + "name": "ggml_mul_mat_skinny_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_skinny_input_f32_bias_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_bias_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_skinny_input_f32_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_residual_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, { "name": "ggml_mul_mat_tiled_input_f32_publish_f32", "family": "loom_libs", diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..1ad784baf5f7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,1459 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.mul_mat.activation_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 2, 5)] + +// Shared contiguous 2D F32 activation skinny matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def pure inline @ggml_q8_lowtoken_geometry(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %q4_format = index.constant 44 : index + %q6_format = index.constant 46 : index + %q8_format = index.constant 9 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %native_q4 = index.cmp eq, %weight_format, %q4_format : index + %native_q6 = index.cmp eq, %weight_format, %q6_format : index + %q8_input = index.cmp eq, %activation_format, %q8_format : index + %q4_min_blocks = index.constant 327680 : index + %blocks = index.div %input_size, %c256 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large_payload = index.cmp uge, %weight_blocks, %q4_min_blocks : index + %q4_can_stage = scalar.andi %native_q4, %large_payload : i1 + %native_staged = scalar.ori %q4_can_stage, %native_q6 : i1 + %native_q8 = scalar.andi %native_staged, %q8_input : i1 + %q4_min_k = index.constant 16384 : index + %q6_min_k = index.constant 8192 : index + %min_k = scf.select %native_q4, %q4_min_k, %q6_min_k : index + %min_n = index.constant 4096 : index + %five_tokens = index.cmp eq, %token_capacity, %c5 : index + %long_k = index.cmp uge, %input_size, %min_k : index + %wide_n = index.cmp uge, %output_size, %min_n : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large = scalar.andi %long_k, %wide_n : i1 + %long_reduction = scalar.andi %large, %contracts : i1 + %staged_shape = scalar.andi %five_tokens, %long_reduction : i1 + %staged = scalar.andi %native_q8, %staged_shape : i1 + %rows = scf.select %native_q4, %c64, %c32 : index + %result = scf.select %staged, %rows, %c0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %block_bytes = index.constant 144 : index + %half_cache_bytes = index.constant 16777216 : index + %cache_bytes = index.constant 33554432 : index + %payload_bytes = index.mul %weight_blocks, %block_bytes : index + %large_half = index.cmp uge, %payload_bytes, %half_cache_bytes : index + %fits_cache = index.cmp ule, %payload_bytes, %cache_bytes : index + %cache_range = scalar.andi %large_half, %fits_cache : i1 + %remainder = index.rem %blocks, %c2 : index + %even_blocks = index.cmp eq, %remainder, %c0 : index + %reduces_rows = index.cmp ult, %output_size, %input_size : index + %native_q4_q8 = scalar.andi %native_q4, %q8_input : i1 + %plain = index.cmp eq, %result, %c0 : index + %q4_plain = scalar.andi %native_q4_q8, %plain : i1 + %five_wide = scalar.andi %five_tokens, %wide_n : i1 + %reduction = scalar.andi %reduces_rows, %even_blocks : i1 + %partition_shape = scalar.andi %five_wide, %reduction : i1 + %partition_payload = scalar.andi %partition_shape, %cache_range : i1 + %partitioned = scalar.andi %q4_plain, %partition_payload : i1 + %k_partitions = scf.select %partitioned, %c2, %c1 : index + func.return %result, %k_partitions : index, index +} + +func.def pure inline @ggml_mul_mat_uses_lowtoken_dot(%token_capacity: index, %input_size: index) -> (i1) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c256 = index.constant 256 : index + %few_tokens = index.cmp ule, %token_capacity, %c5 : index + %tail = index.rem %input_size, %c256 : index + %whole_blocks = index.cmp eq, %tail, %c0 : index + %use_dot = scalar.andi %few_tokens, %whole_blocks : i1 + func.return %use_dot : i1 +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch> pure @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c5 = index.constant 5 : index + %c8 = index.constant 8 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %small_output = index.cmp ult, %output_size, %c64 : index + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%token_capacity, %input_size, %output_size, %weight_format) : (index, index, index, index) -> (index, index) + %staged = index.cmp ugt, %staged_rows, %c0 : index + %staged_rounding = index.sub %staged_rows, %c1 : index + %staged_threads = index.mul %staged_rows, %c16 : index + %cohort_rows = index.div %c8, %k_partitions : index + %cohort_rounding = index.sub %cohort_rows, %c1 : index + %base_row_tile = scf.select %small_output, %c4, %cohort_rows : index + %c1024 = index.constant 1024 : index + %c4096 = index.constant 4096 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %batched_tokens = index.cmp ugt, %token_capacity, %c1 : index + %long_reduction = index.cmp uge, %input_size, %c4096 : index + %partition_tail = index.rem %input_size, %c1024 : index + %full_partitions = index.cmp eq, %partition_tail, %c0 : index + %batch_f32 = scalar.andi %batched_tokens, %f32_input : i1 + %partitionable = scalar.andi %long_reduction, %full_partitions : i1 + %partition_rows = scalar.andi %batch_f32, %partitionable : i1 + %narrow_threads = scf.select %partition_rows, %c1024, %c256 : index + %base_workgroup_size = scf.select %small_output, %narrow_threads, %c128 : index + %base_rounding = scf.select %small_output, %c3, %cohort_rounding : index + %row_rounding = scf.select %staged, %staged_rounding, %base_rounding : index + %row_tile = scf.select %staged, %staged_rows, %base_row_tile : index + %workgroup_size = scf.select %staged, %staged_threads, %base_workgroup_size : index + %rounded_n = index.add %output_size, %row_rounding : index + %direct_tiles = index.div %rounded_n, %row_tile : index + template.return %direct_tiles, %c1, %c1, %workgroup_size : index, index, index, index +} + +func.def inline @ggml_lowtoken_combine_k_partitions(%value: f32, %partitions: index) -> (f32) { + %c1 = index.constant 1 : index + %partitioned = index.cmp ugt, %partitions, %c1 : index + %result = scf.if %partitioned -> (f32) { + %offset = scalar.constant 16 : i32 + %width = scalar.constant 32 : i32 + %peer, %ok = kernel.subgroup.shuffle %value, %offset, %width : f32, i32, i32 + %sum = scalar.addf %value, %peer : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_lowtoken_reduce_cohort_f32(%value: f32) -> (f32) { + %width = scalar.constant 32 : i32 + %off1 = scalar.constant 1 : i32 + %peer1, %ok1 = kernel.subgroup.shuffle %value, %off1, %width : f32, i32, i32 + %sum1 = scalar.addf %value, %peer1 : f32 + %off2 = scalar.constant 2 : i32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %off2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %off4 = scalar.constant 4 : i32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %off4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %off8 = scalar.constant 8 : i32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %off8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + func.return %sum8 : f32 +} + +// Narrow low-token projections assign a full wave to each row. +// Both cohort sizes use the existing wrapper publication hooks. +func.def inline @ggml_lowtoken_reduce_f32(%values: vector<4xf32>, %full_wave: i1) -> (f32) { + %zero = scalar.constant 0.0 : f32 + %width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %sum = vector.reduce %values, %zero : vector<4xf32>, f32 + %peer1, %ok1 = kernel.subgroup.shuffle %sum, %xor1, %width : f32, i32, i32 + %sum1 = scalar.addf %sum, %peer1 : f32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %xor2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %xor4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %xor8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + %result = scf.if %full_wave -> (f32) { + %xor16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %xor16, %width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer16 : f32 + %lane32 = scalar.constant 32 : i32 + %peer32 = kernel.subgroup.broadcast %sum16 from %lane32 : f32, i32 + %sum64 = scalar.addf %sum16, %peer32 : f32 + scf.yield %sum64 : f32 + } else { + scf.yield %sum8 : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %tokens = index.assume %token_count [range(%token_count, 2, 5)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workgroup_size = kernel.workgroup.size : index + %one_wave_per_row = index.cmp eq, %workgroup_size, %c256 : index + %c1024 = index.constant 1024 : index + %partition_group = index.cmp eq, %workgroup_size, %c1024 : index + %narrow_output = index.cmp ult, %n, %c64 : index + %partitioned = scalar.andi %partition_group, %narrow_output : i1 + %full_wave = scalar.ori %one_wave_per_row, %partitioned : i1 + %k_partitions = scf.select %partitioned, %c4, %c1 : index + %cohort_width = scf.select %full_wave, %c64, %c16 : index + %tile_rows = scf.select %full_wave, %c4, %c8 : index + %wave_rows = scf.select %full_wave, %c1, %c4 : index + %step_width = scf.select %full_wave, %c256, %c64 : index + %steps_per_block = scf.select %full_wave, %c1, %c4 : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %lane16 = index.rem %lane, %cohort_width : index + %cohort = index.div %lane, %cohort_width : index + %tile_base = index.mul %channel_tile, %tile_rows : index + %row_subgroup = index.div %subgroup, %k_partitions : index + %k_partition = index.rem %subgroup, %k_partitions : index + %wave_add = index.mul %row_subgroup, %wave_rows : index + %wave_base = index.add %tile_base, %wave_add : index + %row = index.add %wave_base, %cohort : index + %valid_row = index.cmp ult, %row, %n : index + %packet = index.rem %lane16, %c8 : index + %lane_group = index.div %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %blocks = index.div %k, %c256 : index + %steps = index.div %k, %step_width : index + %partition_steps = index.div %steps, %k_partitions : index + %step_begin = index.mul %k_partition, %partition_steps : index + %step_end = index.add %step_begin, %partition_steps : index + %weight_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %row_bytes = index.scale %blocks, %weight_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %a_na, %w_na = buffer.assume.noalias %input, %weight : buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%k]xf32> + %total0, %total1, %total2, %total3, %total4 = scf.for %step = [%step_begin to %step_end step %c1](%acc0 = %zero : vector<4xf32>, %acc1 = %zero : vector<4xf32>, %acc2 = %zero : vector<4xf32>, %acc3 = %zero : vector<4xf32>, %acc4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block = index.div %step, %steps_per_block : index + %quarter = index.rem %step, %steps_per_block : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %step_k = index.mul %step, %step_width : index + %kk = index.add %step_k, %lane_k : index + %ww = scf.if %valid_row -> (vector<4xf32>) { + %format44 = index.constant 44 : index + %packed = index.cmp eq, %weight_format, %format44 : index + %wh = scf.if %packed -> (vector<4xf32>) { + %v = func.call @ggml_q4k_native_row64_f32_vector4(%w_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %v = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %w_na, %row_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } + scf.yield %wh : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next0 = scf.if %valid0 -> (vector<4xf32>) { + %a = vector.load %av[%t0, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc0 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc0 : vector<4xf32> + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next1 = scf.if %valid1 -> (vector<4xf32>) { + %a = vector.load %av[%t1, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc1 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc1 : vector<4xf32> + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next2 = scf.if %valid2 -> (vector<4xf32>) { + %a = vector.load %av[%t2, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc2 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc2 : vector<4xf32> + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next3 = scf.if %valid3 -> (vector<4xf32>) { + %a = vector.load %av[%t3, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc3 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc3 : vector<4xf32> + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next4 = scf.if %valid4 -> (vector<4xf32>) { + %a = vector.load %av[%t4, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc4 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc4 : vector<4xf32> + } + scf.yield %next0, %next1, %next2, %next3, %next4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %c5 = index.constant 5 : index + %stage_width = scf.select %partitioned, %c16, %c8 : index + %stage_elements = index.mul %stage_width, %c5 : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%stage_width]xf32> + %local_row0 = index.add %wave_add, %cohort : index + %partial_offset = index.mul %k_partition, %c4 : index + %local_row = index.add %local_row0, %partial_offset : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total0, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total1, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total2, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total3, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total4, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %tile_rows, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %tile_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, 4), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %stage_width : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_origin = scf.if %partitioned -> (index) { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 64), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } else { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 36), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values0 = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %values = scf.if %partitioned -> (vector<4xf32>) { + %stage_origin1 = index.add %stage_origin, %c4 : index + %stage_origin2 = index.add %stage_origin1, %c4 : index + %stage_origin3 = index.add %stage_origin2, %c4 : index + %values1 = vector.load %flat_stage[%stage_origin1] : view<[%stage_elements]xf32> -> vector<4xf32> + %values2 = vector.load %flat_stage[%stage_origin2] : view<[%stage_elements]xf32> -> vector<4xf32> + %values3 = vector.load %flat_stage[%stage_origin3] : view<[%stage_elements]xf32> -> vector<4xf32> + %sum01 = vector.addf %values0, %values1 : vector<4xf32> + %sum012 = vector.addf %sum01, %values2 : vector<4xf32> + %sum = vector.addf %sum012, %values3 : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values0 : vector<4xf32> + } + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %epilogue, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %row_rounding = index.sub %tile_rows, %c1 : index + %rounded_n = index.add %n, %row_rounding : index + %output_tiles = index.div %rounded_n, %tile_rows : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Copy the original Q8_1_x4 records for up to 2048 reduction elements. +// Call uniformly at a slab boundary; only the final slab may be partial. +func.def inline @ggml_q8_1_x4_stage_lowtoken_slab(%input: buffer, %stage: buffer, %capacity: index, %slab_begin: index, %blocks: index, %global_row_bytes: offset) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c144 = index.constant 144 : index + %c576 = index.constant 576 : index + %c18 = index.constant 18 : index + %base = index.constant 0 : offset + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %workgroup_size = kernel.workgroup.size : index + %thread = kernel.workitem.id : index + %stage_packets = index.mul %capacity, %c144 : index + %stage_elements = index.mul %capacity, %c576 : index + %stage_view = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xi32> + %slab_limit = index.add %slab_begin, %c8 : index + %full_slab = index.cmp ule, %slab_limit, %blocks : index + %slab_end = scf.select %full_slab, %slab_limit, %blocks : index + %active_blocks0 = index.sub %slab_end, %slab_begin : index + %active_blocks = index.assume %active_blocks0 [range(%active_blocks0, 1, 8)] : index + %active_packets = index.mul %active_blocks, %c18 : index + %slab_byte_add = index.scale %slab_begin, %block_bytes : index, offset -> offset + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %copy_packet = [%thread to %stage_packets step %workgroup_size] { + %copy_token = index.div %copy_packet, %c144 : index + %copy_inner = index.rem %copy_packet, %c144 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + scf.if %copy_valid { + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %copy_bytes = index.scale %copy_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slab_byte_add : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + %copy_origin = index.mul %copy_packet, %c4 : index + vector.store %values, %stage_view[%copy_origin] : vector<4xi32>, view<[%stage_elements]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_mul_mat_q8_lowtoken_publish(%tokens: index, %n: index, %row_tile: index, %row_base: index, %local_row: index, %block_lane: index, %channel_tile: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %total_g0: f32, %total_g1: f32, %total_g2: f32, %total_g3: f32, %total_g4: f32, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %leader = index.cmp eq, %block_lane, %c0 : index + %width = scalar.constant 32 : i32 + %stage_rows = index.constant 5 : index + %stage_elements = index.mul %stage_rows, %row_tile : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%row_tile]xf32> + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum = scalar.addf %total_g0, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum = scalar.addf %total_g1, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum = scalar.addf %total_g2, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum = scalar.addf %total_g3, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum = scalar.addf %total_g4, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %row_tile, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %row_limit = index.sub %row_tile, %c4 : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, %row_limit), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %row_tile : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_limit = index.sub %stage_elements, %c4 : index + %stage_origin = index.assume %stage_origin0 [range(%stage_origin0, 0, %stage_limit), mul(%stage_origin0, 4)] : index + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %epilogue, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %rounding = index.sub %row_tile, %c1 : index + %rounded_n = index.add %n, %rounding : index + %output_tiles = index.div %rounded_n, %row_tile : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_q4_prefetch_native_words(%weight: buffer, %row_group_block: index, %block: index, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<2xi32>) { + %zero = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %group_block = index.add %row_group_block, %block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %header_view = buffer.view %weight[%wb] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%wc] : buffer -> view<8x64x4xi32> + %headers = vector.load %header_view[%zero] : view<4xi32> -> vector<4xi32> + %raw = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %headers, %raw : vector<4xi32>, vector<2xi32> +} + +// Native Q4_K bytes and Q8_1_x4 activations, with one 16-lane cohort per row. +// Retain the wrapper publication hooks so bias, residual and RMSNorm fusions +// share this contraction without a separate dispatch or intermediate output. +func.def pure inline @ggml_q4_q8_lowtoken_token_pairs(%capacity: index, %tokens: index, %k: index, %n: index) -> (i1) { + %c4 = index.constant 4 : index + %c4096 = index.constant 4096 : index + %capacity4 = index.cmp eq, %capacity, %c4 : index + %tokens4 = index.cmp eq, %tokens, %c4 : index + %four_tokens = scalar.andi %capacity4, %tokens4 : i1 + %long_input = index.cmp uge, %k, %c4096 : index + %wide_output = index.cmp uge, %n, %c4096 : index + %contracts = index.cmp ule, %n, %k : index + %wide = scalar.andi %long_input, %wide_output : i1 + %contracting = scalar.andi %wide, %contracts : i1 + %eligible = scalar.andi %four_tokens, %contracting : i1 + func.return %eligible : i1 +} + +func.def inline @ggml_q4_q8_lowtoken_adjacent_scale_min(%scale0: i32, %scale1: i32, %scale2: i32, %pair0: index) -> (i32, i32, i32, i32) { + %pair = index.assume %pair0 [range(%pair0, 0, 3)] : index + %two = index.constant 2 : index + %sixteen = index.constant 16 : index + %low = index.cmp ult, %pair, %two : index + %pair_lane = index.rem %pair, %two : index + %shift_index = index.mul %pair_lane, %sixteen : index + %shift = index.cast %shift_index : index to i32 + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %eight_i32 = scalar.constant 8 : i32 + %low_mask = scalar.constant 3855 : i32 + %high_mask = scalar.constant 12336 : i32 + %byte_mask = scalar.constant 255 : i32 + %high_shift = scalar.addi %shift, %two_i32 : i32 + %min_shift = scalar.addi %shift, %four_i32 : i32 + %s_source = scf.select %low, %scale0, %scale2 : i32 + %m_source = scf.select %low, %scale1, %scale2 : i32 + %s_high_shift = scf.select %low, %shift, %high_shift : i32 + %m_low_shift = scf.select %low, %shift, %min_shift : i32 + %s_lo_raw = scalar.shrui %s_source, %shift : i32 + %s_lo = scalar.andi %s_lo_raw, %low_mask : i32 + %s_hi_raw = scalar.shrui %scale0, %s_high_shift : i32 + %s_hi = scalar.andi %s_hi_raw, %high_mask : i32 + %s = scalar.ori %s_lo, %s_hi : i32 + %m_lo_raw = scalar.shrui %m_source, %m_low_shift : i32 + %m_lo = scalar.andi %m_lo_raw, %low_mask : i32 + %m_hi_raw = scalar.shrui %scale1, %s_high_shift : i32 + %m_hi = scalar.andi %m_hi_raw, %high_mask : i32 + %m = scalar.ori %m_lo, %m_hi : i32 + %s0 = scalar.andi %s, %byte_mask : i32 + %m0 = scalar.andi %m, %byte_mask : i32 + %s1 = scalar.shrui %s, %eight_i32 : i32 + %m1 = scalar.shrui %m, %eight_i32 : i32 + func.return %s0, %m0, %s1, %m1 : i32, i32, i32, i32 +} + +func.def inline @ggml_q4_q8_lowtoken_dot8(%paired_tokens: i1, %weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %value = scf.if %paired_tokens -> (f32) { + %zero = vector.constant 0 : vector<1xi32> + %wi0 = vector.extract %weight[0] : vector<2xi32> -> i32 + %ai0 = vector.extract %activation[0] : vector<2xi32> -> i32 + %ww0 = vector.splat %wi0 : vector<1xi32> + %aw0 = vector.splat %ai0 : vector<1xi32> + %w0 = vector.bitcast %ww0 : vector<1xi32> to vector<4xi8> + %a0 = vector.bitcast %aw0 : vector<1xi32> to vector<4xi8> + %d0 = vector.dot4i %w0, %a0, %zero : vector<4xi8>, vector<4xi8>, vector<1xi32> + %wi1 = vector.extract %weight[1] : vector<2xi32> -> i32 + %ai1 = vector.extract %activation[1] : vector<2xi32> -> i32 + %ww1 = vector.splat %wi1 : vector<1xi32> + %aw1 = vector.splat %ai1 : vector<1xi32> + %w1 = vector.bitcast %ww1 : vector<1xi32> to vector<4xi8> + %a1 = vector.bitcast %aw1 : vector<1xi32> to vector<4xi8> + %d1 = vector.dot4i %w1, %a1, %d0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %sum = vector.extract %d1[0] : vector<1xi32> -> i32 + %result = scalar.sitofp %sum : i32 to f32 + scf.yield %result : f32 + } else { + %original = func.call @ggml_dot_u8_s8_vector8_f32(%weight, %activation) : (vector<2xi32>, vector<2xi32>) -> (f32) + scf.yield %original : f32 + } + func.return %value : f32 +} + +func.def inline @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %paired_tokens = func.call pure @ggml_q4_q8_lowtoken_token_pairs(%capacity, %tokens, %k, %n) : (index, index, index, index) -> (i1) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = index.add %channel_tile, %c0 : index + %c64_rows = index.constant 64 : index + %direct_rows = index.div %c8, %k_partitions : index + %row_tile = scf.select %staged, %c64_rows, %direct_rows : index + %row_base = index.mul %wg, %row_tile : index + %wave_rows = index.div %c4, %k_partitions : index + %wave_row = index.mul %wave, %wave_rows : index + %cohort_width = index.mul %c16, %k_partitions : index + %row_half = index.div %lane, %cohort_width : index + %cohort_lane = index.rem %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %k_partitions : index + %k_team = index.div %cohort_lane, %c16 : index + %k_begin = index.mul %k_team, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %global_q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %slab_row_bytes = index.constant 2304 : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %channel_tile, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = index.constant 8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_header = vector.constant 0 : vector<4xi32> + %zero_raw = vector.constant 0 : vector<2xi32> + %prime_header, %prime_raw = scf.if %staged -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %c0, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %zero_header, %zero_raw : vector<4xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_header, %last_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_header = %prime_header : vector<4xi32>, %carry_raw = %prime_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32>) { + %block_in_slab = index.rem %block_base, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %new_slab = scalar.andi %staged, %slab_boundary : i1 + scf.if %new_slab { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block_base, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %block = index.add %block_base, %k_begin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_block, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %staged -> (vector<4xi32>) { + scf.yield %carry_header : vector<4xi32> + } else { + %header = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %header : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %pair_scale0, %pair_min0, %pair_scale1, %pair_min1 = scf.if %paired_tokens -> (i32, i32, i32, i32) { + %s0, %m0, %s1, %m1 = func.call @ggml_q4_q8_lowtoken_adjacent_scale_min(%g_s0, %g_s1, %g_s2, %pair) : (i32, i32, i32, index) -> (i32, i32, i32, i32) + scf.yield %s0, %m0, %s1, %m1 : i32, i32, i32, i32 + } else { + %zero_i32 = scalar.constant 0 : i32 + scf.yield %zero_i32, %zero_i32, %zero_i32, %zero_i32 : i32, i32, i32, i32 + } + %g_scale0i, %g_min0i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale0, %pair_min0 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale1, %pair_min1 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %staged -> (vector<2xi32>) { + scf.yield %carry_raw : vector<2xi32> + } else { + %raw = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %raw : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_header, %next_raw = scf.if %has_next -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %next_block, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %carry_header, %carry_raw : vector<4xi32>, vector<2xi32> + } + %input_block = scf.select %staged, %block_in_slab, %safe_block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0 = scf.if %valid0 -> (f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ad0h = view.load %ds_view[%q8_meta0] : view<8xf16> -> f16 + %ad1h = view.load %ds_view[%q8_meta1] : view<8xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = view.load %ds_view[%q8_sum0] : view<8xf16> -> f16 + %as1h = view.load %ds_view[%q8_sum1] : view<8xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1 = scf.if %valid1 -> (f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + scf.if %paired_tokens { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2 = scf.if %valid2 -> (f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3 = scf.if %valid3 -> (f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4 = scf.if %valid4 -> (f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_header, %next_raw : f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32> + } + %combined0 = func.call @ggml_lowtoken_combine_k_partitions(%total_g0, %k_partitions) : (f32, index) -> (f32) + %combined1 = func.call @ggml_lowtoken_combine_k_partitions(%total_g1, %k_partitions) : (f32, index) -> (f32) + %combined2 = func.call @ggml_lowtoken_combine_k_partitions(%total_g2, %k_partitions) : (f32, index) -> (f32) + %combined3 = func.call @ggml_lowtoken_combine_k_partitions(%total_g3, %k_partitions) : (f32, index) -> (f32) + %combined4 = func.call @ggml_lowtoken_combine_k_partitions(%total_g4, %k_partitions) : (f32, index) -> (f32) + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %cohort_lane, %channel_tile, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %combined0, %combined1, %combined2, %combined3, %combined4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Native Q6_K carrier with signed I8 dot products and F32 scaling. +func.def inline @ggml_q6_prefetch_native_words(%weight_na: buffer, %row_group_block: index, %block: index, %row_lane: index, %half16: index, %group: index, %half128: index, %ql_side: index, %packet: index) -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %c2 = index.constant 2 : index + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %group_block = index.add %row_group_block, %block : index + %block_base = index.scale %group_block, %record_bytes : index, offset -> offset + %scale_base = index.add %block_base, %scale_add : offset + %raw_base = index.add %block_base, %raw_add : offset + %d_view = buffer.view %weight_na[%block_base] : buffer -> view<64xf16> + %s_view = buffer.view %weight_na[%scale_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_na[%raw_base] : buffer -> view<2x3x64x8xi32> + %dh = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %si = view.load %s_view[%row_lane, %half16, %group] : view<64x2x8xi8> -> i8 + %ql = vector.load %raw_view[%half128, %ql_side, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + %qh = vector.load %raw_view[%half128, %c2, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + func.return %dh, %si, %ql, %qh : f16, i8, vector<4xi32>, vector<4xi32> +} + +func.def inline @ggml_q6_q8_lowtoken_activation_values(%pipelined: i1, %token: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index) -> (vector<16xi8>, f32) { + %code_offset = index.constant 16 : offset + %a, %ad = scf.if %pipelined -> (vector<16xi8>, f32) { + %base = index.constant 0 : offset + %c4 = index.constant 4 : index + %q8_block = index.assume %q8_block0 [range(%q8_block0, 0, 7)] : index + %word = index.add %q8_word0, %c4 : index + %q_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x36xi32> + %ds_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x72xf16> + %words = vector.load %q_view[%current_slot, %token, %q8_block, %word] : view<2x[%capacity]x8x36xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%current_slot, %token, %q8_block, %q8_meta] : view<2x[%capacity]x8x72xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } else { + %tb = index.scale %token, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %words = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%q8_meta] : view<8xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } + func.return %a, %ad : vector<16xi8>, f32 +} + +func.def inline @ggml_q6_q8_lowtoken_dot16(%pipelined: i1, %w: vector<16xi8>, %a: vector<16xi8>) -> (i32) { + %sum = scf.if %pipelined -> (i32) { + %di0 = vector.constant 0 : vector<1xi32> + %dw0 = vector.slice %w[0] : vector<16xi8> -> vector<4xi8> + %da0 = vector.slice %a[0] : vector<16xi8> -> vector<4xi8> + %di1 = vector.dot4i %dw0, %da0, %di0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw1 = vector.slice %w[4] : vector<16xi8> -> vector<4xi8> + %da1 = vector.slice %a[4] : vector<16xi8> -> vector<4xi8> + %di2 = vector.dot4i %dw1, %da1, %di1 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw2 = vector.slice %w[8] : vector<16xi8> -> vector<4xi8> + %da2 = vector.slice %a[8] : vector<16xi8> -> vector<4xi8> + %di3 = vector.dot4i %dw2, %da2, %di2 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw3 = vector.slice %w[12] : vector<16xi8> -> vector<4xi8> + %da3 = vector.slice %a[12] : vector<16xi8> -> vector<4xi8> + %di4 = vector.dot4i %dw3, %da3, %di3 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %chain_sum = vector.extract %di4[0] : vector<1xi32> -> i32 + scf.yield %chain_sum : i32 + } else { + %init = vector.constant 0 : vector<4xi32> + %dot = vector.dot4i %w, %a, %init : vector<16xi8>, vector<16xi8>, vector<4xi32> + %seed = scalar.constant 0 : i32 + %direct_sum = vector.reduce %dot, %seed : vector<4xi32>, i32 + scf.yield %direct_sum : i32 + } + func.return %sum : i32 +} + +func.def inline @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input: buffer, %capacity0: index, %slot_begin: index, %blocks: index, %global_row_bytes: offset) -> (vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c18 = index.constant 18 : index + %c72 = index.constant 72 : index + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %copy_packet = index.rem %thread, %total_packets : index + %copy_token = index.div %copy_packet, %c72 : index + %copy_inner = index.rem %copy_packet, %c72 : index + %remaining = index.sub %blocks, %slot_begin : index + %full = index.cmp uge, %remaining, %c4 : index + %active0 = scf.select %full, %c4, %remaining : index + %active = index.assume %active0 [range(%active0, 1, 4)] : index + %active_packets = index.mul %active, %c18 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + %safe_inner = scf.select %copy_valid, %copy_inner, %c0 : index + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %slot_bytes = index.scale %slot_begin, %block_bytes : index, offset -> offset + %copy_bytes = index.scale %safe_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slot_bytes : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + func.return %values : vector<4xi32> +} +func.def inline @ggml_q8_1_x4_publish_lowtoken_half_slab(%stage: buffer, %capacity0: index, %slot0: index, %values: vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %slot = index.assume %slot0 [range(%slot0, 0, 1)] : index + %c4 = index.constant 4 : index + %c72 = index.constant 72 : index + %base = index.constant 0 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %thread_valid = index.cmp ult, %thread, %total_packets : index + scf.if %thread_valid { + %copy_token0 = index.div %thread, %c72 : index + %copy_token = index.assume %copy_token0 [range(%copy_token0, 0, 4), lt(%copy_token0, %capacity)] : index + %copy_inner = index.rem %thread, %c72 : index + %origin = index.mul %copy_inner, %c4 : index + %destination = buffer.view %stage[%base] : buffer -> view<2x[%capacity]x288xi32> + vector.store %values, %destination[%slot, %copy_token, %origin] : vector<4xi32>, view<2x[%capacity]x288xi32> + } + func.return +} + + +func.def inline @ggml_q6_q8_lowtoken_accumulate(%pipelined: i1, %tokens: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index, %w: vector<16xi8>, %ws: f32, %acc_g0: f32, %acc_g1: f32, %acc_g2: f32, %acc_g3: f32, %acc_g4: f32) -> (f32, f32, f32, f32, f32) { + %code_offset = index.constant 16 : offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %computed_g0 = scf.if %valid0 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t0, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g0 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %computed_g1 = scf.if %valid1 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t1, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g1 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %computed_g2 = scf.if %valid2 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t2, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g2 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %computed_g3 = scf.if %valid3 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t3, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g3 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %computed_g4 = scf.if %valid4 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t4, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g4 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + func.return %computed_g0, %computed_g1, %computed_g2, %computed_g3, %computed_g4 : f32, f32, f32, f32, f32 +} + +func.def inline @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %row_tile = scf.select %staged, %c32, %c8 : index + %row_base = index.mul %channel_tile, %row_tile : index + %wave_row = index.mul %wave, %c4 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %block_lane = index.rem %lane, %c16 : index + %group = index.div %block_lane, %c2 : index + %half16 = index.rem %block_lane, %c2 : index + %half128 = index.div %group, %c4 : index + %group_in_half = index.rem %group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift0 = index.mul %nibble, %c4 : index + %nibble_shift1 = index.cast %nibble_shift0 : index to i32 + %nibble_shift = vector.splat %nibble_shift1 : vector<4xi32> + %high_shift0 = index.mul %group_in_half, %c2 : index + %high_shift1 = index.cast %high_shift0 : index to i32 + %high_shift = vector.splat %high_shift1 : vector<4xi32> + %packet = index.mul %half16, %c4 : index + %q8_group_words = index.mul %group_in_half, %c8 : index + %q8_word0 = index.add %q8_group_words, %packet : index + %q8_meta = index.mul %group_in_half, %c2 : index + %blocks = index.div %k, %c256 : index + %q8_groups = index.div %k, %c128 : index + %q8_block_bytes = index.constant 144 : offset + %global_q8_row_bytes = index.scale %q8_groups, %q8_block_bytes : index, offset -> offset + %native_block_bytes = index.constant 210 : index + %cache_bytes = index.constant 33554432 : index + %weight_blocks = index.mul %blocks, %n : index + %weight_bytes = index.mul %weight_blocks, %native_block_bytes : index + %exceeds_cache = index.cmp uge, %weight_bytes, %cache_bytes : index + %pipelined = scalar.andi %staged, %exceeds_cache : i1 + %full_slab_row_bytes = index.constant 2304 : offset + %half_slab_row_bytes = index.constant 1152 : offset + %slab_row_bytes = scf.select %pipelined, %half_slab_row_bytes, %full_slab_row_bytes : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %code_offset = index.constant 16 : offset + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %staged_row_group = index.div %channel_tile, %c2 : index + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %staged_row_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %input_na, %weight_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = scf.select %pipelined, %c4, %c8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_half = scalar.constant 0.0 : f16 + %zero_code = scalar.constant 0 : i8 + %zero_words = vector.constant 0 : vector<4xi32> + %prime_d, %prime_s, %prime_ql, %prime_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %pd, %ps, %pql, %pqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %c0, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %pd, %ps, %pql, %pqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %zero_half, %zero_code, %zero_words, %zero_words : f16, i8, vector<4xi32>, vector<4xi32> + } + scf.if %pipelined { + %prime_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %c0, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %c0, %prime_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_d, %last_s, %last_ql, %last_qh = scf.for %block = [%c0 to %blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_d = %prime_d : f16, %carry_s = %prime_s : i8, %carry_ql = %prime_ql : vector<4xi32>, %carry_qh = %prime_qh : vector<4xi32>) -> (f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32>) { + %block_in_slab = index.rem %block, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %direct_stage = scalar.andi %staged, %slab_boundary : i1 + %true = scalar.constant true : i1 + %not_pipelined = scalar.xori %pipelined, %true : i1 + %stage_current = scalar.andi %direct_stage, %not_pipelined : i1 + scf.if %stage_current { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %slab_number = index.div %block, %slab_blocks : index + %current_slot = index.rem %slab_number, %c2 : index + %following_slab = index.add %slab_number, %c1 : index + %following_begin = index.mul %following_slab, %slab_blocks : index + %following_exists = index.cmp ult, %following_begin, %blocks : index + %next_slot = index.rem %following_slab, %c2 : index + %staged_future = scalar.andi %pipelined, %following_exists : i1 + %current_d, %current_s, %current_ql, %current_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } + %d = scalar.extf %current_d : f16 to f32 + %sf = scalar.sitofp %current_s : i8 to f32 + %ws = scalar.mulf %d, %sf : f32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %four = vector.constant 4 : vector<4xi32> + %ql_shifted = vector.shrui %current_ql, %nibble_shift : vector<4xi32> + %low = vector.andi %ql_shifted, %nibble_mask : vector<4xi32> + %qh_shifted = vector.shrui %current_qh, %high_shift : vector<4xi32> + %high0 = vector.andi %qh_shifted, %high_mask : vector<4xi32> + %high = vector.shli %high0, %four : vector<4xi32> + %code = vector.ori %low, %high : vector<4xi32> + // Guard bits prevent subtraction from borrowing across packed bytes. + %guard_mask = vector.constant -2139062144 : vector<4xi32> + %center = vector.constant 538976288 : vector<4xi32> + %guarded = vector.ori %code, %guard_mask : vector<4xi32> + %centered = vector.subi %guarded, %center : vector<4xi32> + %signed_code = vector.xori %centered, %guard_mask : vector<4xi32> + %w = vector.bitcast %signed_code : vector<4xi32> to vector<16xi8> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_d, %next_s, %next_ql, %next_qh = scf.if %has_next -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %next_block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } + %input_block = scf.select %staged, %block_in_slab, %block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %half128 : index + %q8_local_block_add = index.scale %q8_block, %q8_block_bytes : index, offset -> offset + %slot_row = index.mul %current_slot, %capacity : index + %slot_byte_add0 = index.scale %slot_row, %slab_row_bytes : index, offset -> offset + %slot_byte_add = scf.select %pipelined, %slot_byte_add0, %base : offset + %q8_block_add = index.add %slot_byte_add, %q8_local_block_add : offset + %slot_complete = index.cmp eq, %block_in_slab, %c3 : index + %publish_next = scalar.andi %staged_future, %slot_complete : i1 + %next_g0, %next_g1, %next_g2, %next_g3, %next_g4 = scf.if %publish_next -> (f32, f32, f32, f32, f32) { + %next_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %following_begin, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %next_slot, %next_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } else { + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_d, %next_s, %next_ql, %next_qh : f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32> + } + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %block_lane, %channel_tile, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %activation_format = config.get @ggml.mul_mat.activation_format : index + %q8_format = index.constant 9 : index + %uses_q8 = index.cmp eq, %activation_format, %q8_format : index + scf.if %uses_q8 { + %q6_format = index.constant 46 : index + %uses_q6 = index.cmp eq, %weight_format, %q6_format : index + scf.if %uses_q6 { + func.call @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + func.call @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token2(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 2 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token3(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 3 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token4(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 4 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token5(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 5 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.assume %token_count [range(%token_count, 2, 5)] : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %is_token2 = index.cmp eq, %tokens, %c2 : index + scf.if %is_token2 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token2(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + %is_token3 = index.cmp eq, %tokens, %c3 : index + scf.if %is_token3 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token3(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + %is_token4 = index.cmp eq, %tokens, %c4 : index + scf.if %is_token4 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token4(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token5(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom new file mode 100644 index 000000000000..24d2078df656 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom @@ -0,0 +1,31 @@ +template.decl @ggml.matmul.skinny_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.def<@ggml.matmul.skinny_postprocess_f32.add_bias_vector4> device @ggml_matmul_skinny_postprocess_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.skinny_postprocess_f32.add_residual_vector4> device @ggml_matmul_skinny_postprocess_f32_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom new file mode 100644 index 000000000000..519714efb716 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom @@ -0,0 +1,14 @@ +template.decl @ggml.matmul.skinny_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.skinny_publish_f32.store_vector4> device @ggml_matmul_skinny_publish_f32_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom new file mode 100644 index 000000000000..5135b662819c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom @@ -0,0 +1,188 @@ +// Skinny GGML matmul for contiguous 2D F32 activations. +// +// The accumulation core covers low-token direct F32 and staged quantized +// activation paths. This wrapper selects the output epilogue. +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 2, 5)] + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_skinny_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 2, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %accumulated = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.skinny_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %accumulated, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %accumulated : vector<4xf32> + } + %residual_added = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.skinny_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + %published = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %residual_added) : (index, vector<4xf32>) -> (vector<4xf32>) + scf.if %uses_residual { + template.apply<@ggml.matmul.skinny_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.skinny_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + template.return +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile> device @ggml_matmul_skinny_input_f32_publish_f32_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat.input_size : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_bias_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 1 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 2 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = index.constant 3 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom index da096db9f721..9e7cde4571b4 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom @@ -32,7 +32,7 @@ config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, config.decl @ggml.mul_mat.activation_format : %value: index where [range(%value, 0, 16)] -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_wmma.loom deleted file mode 100644 index 95cd5db63032..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_wmma.loom +++ /dev/null @@ -1,55 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) - -template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) - -amdgpu.target @ggml_mul_mat_add_gfx11_wave64 {subgroup_size = 64} - -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] - -config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_add_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - scf.if %publish_word { - %published = template.apply<@ggml.mul_mat_f32_f32_wmma.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %values, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - template.apply<@ggml.mul_mat_f32_f32_wmma.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_add_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.return -} - -kernel.def target(@ggml_mul_mat_add_gfx11_wave64) @ggml_mul_mat_add_f32_f32_wmma(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom deleted file mode 100644 index cd630d060222..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_next_rmsnorm_f32_f32_wmma.loom +++ /dev/null @@ -1,68 +0,0 @@ -// Dense GGML matmul fused with the following residual add and RMSNorm-weight -// multiply: -// MUL_MAT -> ADD(residual) -> RMS_NORM -> MUL(norm_weight) -// -// The residual-add output is explicitly materialized for later graph users. -// The raw matmul and raw RMSNorm outputs are internal to this fused dispatch. -template.decl @ggml.mul_mat_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: index, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) - -template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) - -amdgpu.target @ggml_mul_mat_add_next_rmsnorm_gfx11_wave64 {subgroup_size = 64} - -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @ggml.mul_mat_postops.rms_epsilon : f32 - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_add_next_rmsnorm_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index - %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index - scf.if %publish_word { - %published = template.apply<@ggml.mul_mat_f32_f32_wmma.add_residual_vector4>(%token_count, %output_size, %token0, %channel, %values, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - template.apply<@ggml.mul_mat_f32_f32_wmma.store_vector4>(%token_count, %output_size, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_add_next_rmsnorm_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.apply<@ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight>(%token_count, %output_size, %output_tile_count, %token_tile_count, %epsilon, %token_tile, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, f32, index, buffer, buffer, buffer, buffer) - template.return -} - -kernel.def target(@ggml_mul_mat_add_next_rmsnorm_gfx11_wave64) @ggml_mul_mat_add_next_rmsnorm_f32_f32_wmma(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %epsilon = config.get @ggml.mul_mat_postops.rms_epsilon : f32 - %c0 = index.constant 0 : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epsilon, %channel_tile, %token_tile, %input, %weight, %residual_output, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_f32_f32_wmma.loom deleted file mode 100644 index 27c9fef60c0e..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_f32_f32_wmma.loom +++ /dev/null @@ -1,58 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_wmma.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) - -template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) - -amdgpu.target @ggml_mul_mat_bias_add_gfx11_wave64 {subgroup_size = 64} - -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] - -config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_bias_add_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - scf.if %publish_word { - %biased = template.apply<@ggml.mul_mat_f32_f32_wmma.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - %published = template.apply<@ggml.mul_mat_f32_f32_wmma.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - template.apply<@ggml.mul_mat_f32_f32_wmma.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_bias_add_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.return -} - -kernel.def target(@ggml_mul_mat_bias_add_gfx11_wave64) @ggml_mul_mat_bias_add_f32_f32_wmma(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom deleted file mode 100644 index aa6b4bfd4148..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_add_next_rmsnorm_f32_f32_wmma.loom +++ /dev/null @@ -1,63 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_wmma.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: index, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) - -template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) - -amdgpu.target @ggml_mul_mat_bias_add_next_rmsnorm_gfx11_wave64 {subgroup_size = 64} - -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @ggml.mul_mat_postops.rms_epsilon : f32 - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_bias_add_next_rmsnorm_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - scf.if %publish_word { - %biased = template.apply<@ggml.mul_mat_f32_f32_wmma.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - %published = template.apply<@ggml.mul_mat_f32_f32_wmma.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - template.apply<@ggml.mul_mat_f32_f32_wmma.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_bias_add_next_rmsnorm_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.apply<@ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight>(%token_count, %output_size, %output_tile_count, %token_tile_count, %epsilon, %token_tile, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, f32, index, buffer, buffer, buffer, buffer) - template.return -} - -kernel.def target(@ggml_mul_mat_bias_add_next_rmsnorm_gfx11_wave64) @ggml_mul_mat_bias_add_next_rmsnorm_f32_f32_wmma(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %epsilon = config.get @ggml.mul_mat_postops.rms_epsilon : f32 - %c0 = index.constant 0 : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_f32_f32_wmma.loom deleted file mode 100644 index cac1acce67bf..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_bias_f32_f32_wmma.loom +++ /dev/null @@ -1,55 +0,0 @@ -template.decl @ggml.mul_mat_f32_f32_wmma.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) - -template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) - -template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) - -template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) - -amdgpu.target @ggml_mul_mat_bias_gfx11_wave64 {subgroup_size = 64} - -config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] - -config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] - -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_bias_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - scf.if %publish_word { - %biased = template.apply<@ggml.mul_mat_f32_f32_wmma.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) - template.apply<@ggml.mul_mat_f32_f32_wmma.store_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %output) : (index, index, index, index, vector<4xf32>, buffer) - } - template.return -} - -template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_bias_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - template.return -} - -kernel.def target(@ggml_mul_mat_bias_gfx11_wave64) @ggml_mul_mat_bias_f32_f32_wmma(%token_count: index) { - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %c1 = index.constant 1 : index - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer) { - %input_size = config.get @ggml.mul_mat_postops.input_size : index - %output_size = config.get @ggml.mul_mat_postops.output_size : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %weight_format = config.get @ggml.mul_mat_postops.weight_format : index - %c0 = index.constant 0 : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom index 209d2650725a..6e7f8899cdfb 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom @@ -181,88 +181,6 @@ kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_q4_k_f16_wmma_prefil kernel.return } -kernel.def target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_narrow_split_k4(%token_count: index) { - %token_capacity = config.get @ggml.workload.token_capacity : index - %one = index.constant 1 : index - %split_count = index.constant 4 : index - %tile_tokens = index.constant 32 : index - %workgroup_size = index.constant 128 : index - %token_tiles = index.div %token_capacity, %tile_tokens : index - kernel.launch.config workgroups(%one, %token_tiles, %split_count) workgroup_size(%workgroup_size, %one, %one) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %final_output: buffer, %counters: buffer) { - %input_size0 = config.get @ggml.mul_mat.input_size : index - %input_size = index.assume %input_size0 [range(%input_size0, 4096, 32768), mul(%input_size0, 1024)] : index - %output_size0 = config.get @ggml.mul_mat.output_size : index - %output_size = index.assume %output_size0 [range(%output_size0, 4, 64), mul(%output_size0, 4)] : index - %token_capacity = config.get @ggml.workload.token_capacity : index - %output_accumulation = index.constant 0 : index - %output_unary_op = index.constant 23 : index - %weight_format = config.get @ggml.mul_mat.weight_format : index - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32), le(%token_count, %token_capacity), le(%token_capacity, %token_count)] : index - %channel_tile = kernel.workgroup.id : index - %split_count = index.constant 4 : index - %token_tile = kernel.workgroup.id : index - template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - %base = index.constant 0 : offset - %zero = index.constant 0 : index - %four = index.constant 4 : index - %tile_m = index.constant 32 : index - %step = index.constant 512 : index - %zero_i32 = scalar.constant 0 : i32 - %one_i32 = scalar.constant 1 : i32 - %three_i32 = scalar.constant 3 : i32 - %negative_four_i32 = scalar.constant -4 : i32 - %tid = kernel.workitem.id : index - %first_thread = index.cmp eq, %tid, %zero : index - %counter_count = index.div %token_capacity, %tile_m : index - %counter_view = buffer.view %counters[%base] : buffer -> view<[%counter_count]xi32> - kernel.barrier scope(workgroup) ordering(release) - %arrival_local = scf.if %first_thread -> (i32) { - %previous = view.atomic.rmw %one_i32, %counter_view[%token_tile] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 - scf.yield %previous : i32 - } else { - scf.yield %zero_i32 : i32 - } - %arrival = kernel.workgroup.reduce %arrival_local : i32 - %last = scalar.cmpi eq, %arrival, %three_i32 : i32 - scf.if %last { - kernel.barrier scope(workgroup) ordering(acquire) - %byte_size = index.constant 4 : offset - %plane_elements = index.mul %bounded_token_count, %output_size : index - %plane_bytes = index.scale %plane_elements, %byte_size : index, offset -> offset - %plane2 = index.add %plane_bytes, %plane_bytes : offset - %plane3 = index.add %plane2, %plane_bytes : offset - %partial0 = buffer.view %output[%base] : buffer -> view<[%bounded_token_count]x[%output_size]xf32> - %partial1 = buffer.view %output[%plane_bytes] : buffer -> view<[%bounded_token_count]x[%output_size]xf32> - %partial2 = buffer.view %output[%plane2] : buffer -> view<[%bounded_token_count]x[%output_size]xf32> - %partial3 = buffer.view %output[%plane3] : buffer -> view<[%bounded_token_count]x[%output_size]xf32> - %final_view = buffer.view %final_output[%base] : buffer -> view<[%bounded_token_count]x[%output_size]xf32> - %token_origin = index.mul %token_tile, %tile_m : index - %elements = index.mul %tile_m, %output_size : index - %first_element = index.mul %tid, %four : index - scf.for %linear = [%first_element to %elements step %step] { - %row = index.div %linear, %output_size : index - %channel = index.rem %linear, %output_size : index - %token = index.add %token_origin, %row : index - %v0 = vector.load %partial0[%token, %channel] : view<[%bounded_token_count]x[%output_size]xf32> -> vector<4xf32> - %v1 = vector.load %partial1[%token, %channel] : view<[%bounded_token_count]x[%output_size]xf32> -> vector<4xf32> - %v2 = vector.load %partial2[%token, %channel] : view<[%bounded_token_count]x[%output_size]xf32> -> vector<4xf32> - %v3 = vector.load %partial3[%token, %channel] : view<[%bounded_token_count]x[%output_size]xf32> -> vector<4xf32> - %sum01 = vector.addf %v0, %v1 : vector<4xf32> - %sum012 = vector.addf %sum01, %v2 : vector<4xf32> - %sum = vector.addf %sum012, %v3 : vector<4xf32> - vector.store %sum, %final_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%output_size]xf32> - scf.yield - } - kernel.barrier scope(workgroup) ordering(release) - scf.if %first_thread { - view.atomic.reduce %negative_four_i32, %counter_view[%token_tile] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> - } - } - kernel.return -} - func.decl @ggml_mul_mat_quantized_f16_prefill_conv4(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %weight: buffer, %state: buffer, %filter: buffer, %output: buffer, %cache: buffer) kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_quantized_f16_wmma_prefill_conv4(%token_count: index) { diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 56d5552c1f25..26045fba9cd0 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -11,6 +11,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -150,6 +153,11 @@ static bool environment_flag_enabled(const char * name) { return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0'); } +static const char * environment_string_value(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0') ? value : nullptr; +} + static size_t environment_size_value(const char * name, size_t fallback) { const char * value = std::getenv(name); if (value == nullptr || value[0] == '\0') { @@ -293,6 +301,208 @@ static bool resource_access_writes(ResourceAccess access) { return access == ResourceAccess::Write || access == ResourceAccess::ReadWrite; } +static bool resource_access_reads(ResourceAccess access) { + return access == ResourceAccess::Read || access == ResourceAccess::ReadWrite; +} + +static bool dispatch_trace_enabled() { + return environment_string_value("GGML_HRX_TRACE_DISPATCH_DIR") != nullptr; +} + +static bool dispatch_trace_blob_capture_enabled() { + return environment_flag_enabled("GGML_HRX_TRACE_BLOBS"); +} + +static bool dispatch_trace_should_trace_program(size_t program_command_count) { + const char * value = environment_string_value("GGML_HRX_TRACE_PROGRAM_COMMANDS"); + return value == nullptr || program_command_count == std::strtoull(value, nullptr, 0); +} + +static bool dispatch_trace_should_trace_command(size_t program_command_count, size_t command_index) { + if (!dispatch_trace_should_trace_program(program_command_count)) { + return false; + } + const size_t start = environment_size_value("GGML_HRX_TRACE_START_COMMAND", 0); + const size_t end = + environment_size_value("GGML_HRX_TRACE_END_COMMAND", std::numeric_limits::max()); + return command_index >= start && command_index <= end; +} + +static uint64_t dispatch_trace_hash_bytes(const uint8_t * data, size_t size) { + uint64_t hash = UINT64_C(1469598103934665603); + for (size_t i = 0; i < size; ++i) { + hash ^= static_cast(data[i]); + hash *= UINT64_C(1099511628211); + } + return hash; +} + +static std::string dispatch_trace_hex(uint64_t value) { + std::ostringstream out; + out << std::hex << std::setw(16) << std::setfill('0') << value; + return out.str(); +} + +static std::string dispatch_trace_json_escape(const std::string & text) { + std::string escaped; + escaped.reserve(text.size()); + for (const unsigned char c : text) { + switch (c) { + case '\\': escaped += "\\\\"; break; + case '"': escaped += "\\\""; break; + case '\b': escaped += "\\b"; break; + case '\f': escaped += "\\f"; break; + case '\n': escaped += "\\n"; break; + case '\r': escaped += "\\r"; break; + case '\t': escaped += "\\t"; break; + default: + if (c < 0x20) { + static constexpr char kHex[] = "0123456789abcdef"; + escaped += "\\u00"; + escaped.push_back(kHex[c >> 4]); + escaped.push_back(kHex[c & 0xf]); + } else { + escaped.push_back(static_cast(c)); + } + break; + } + } + return escaped; +} + +static std::string dispatch_trace_kernel_name(const CommandProgramExecutionContext & context, + const PreparedCommand & command) { + if (context.corpus != nullptr && context.target != nullptr) { + const KernelResolveResult resolved = + resolve_kernel_definition(*context.corpus, context.target, command.kernel.specialization.kernel_id); + if (resolved.found()) { + return kernel_definition_name(*resolved.definition); + } + } + std::ostringstream out; + out << "kernel_id=" << command.kernel.specialization.kernel_id; + return out.str(); +} + +static Status dispatch_trace_download_binding(const CommandProgramExecutionContext & context, + const PreparedCommandBinding & binding, + std::vector & data) { + Status status; + data.resize(binding.ref.length); + if (binding.ref.length == 0) { + return status; + } + if (context.host_transfers == nullptr) { + status.log("missing HRX host transfer manager for dispatch trace"); + return status; + } + status.append(context.host_transfers->download_synchronous(context.stream, binding.ref.buffer, binding.ref.offset, + data.data(), binding.ref.length)); + return status; +} + +static Status dispatch_trace_write_blob(const std::filesystem::path & path, const std::vector & data) { + Status status; + try { + std::filesystem::create_directories(path.parent_path()); + std::ofstream out(path, std::ios::binary | std::ios::trunc); + if (!out) { + status.log("failed to create dispatch trace blob %s", path.string().c_str()); + return status; + } + out.write(reinterpret_cast(data.data()), static_cast(data.size())); + } catch (const std::exception & error) { + status.log("failed to write dispatch trace blob %s: %s", path.string().c_str(), error.what()); + } + return status; +} + +static Status dispatch_trace_command_bindings(const CommandProgramExecutionContext & context, + const PreparedCommand & command, + const char * list_kind, + size_t list_size, + size_t command_index, + size_t program_command_count, + const char * stage) { + Status status; + const char * directory_value = environment_string_value("GGML_HRX_TRACE_DISPATCH_DIR"); + if (directory_value == nullptr || !dispatch_trace_should_trace_command(program_command_count, command_index)) { + return status; + } + const bool before_stage = std::strcmp(stage, "before") == 0; + const bool after_stage = std::strcmp(stage, "after") == 0; + if (!before_stage && !after_stage) { + return status; + } + const bool capture_blobs = dispatch_trace_blob_capture_enabled(); + if (before_stage && !capture_blobs) { + return status; + } + const std::filesystem::path root(directory_value); + const std::filesystem::path blobs = root / "blobs"; + const std::filesystem::path trace_path = root / "trace.jsonl"; + const std::string kernel_name = dispatch_trace_kernel_name(context, command); + for (size_t binding_index = 0; binding_index < command.kernel.bindings.size(); ++binding_index) { + const PreparedCommandBinding & binding = command.kernel.bindings[binding_index]; + const bool should_capture = before_stage ? resource_access_reads(binding.binding.access) : + resource_access_writes(binding.binding.access); + if (!should_capture) { + continue; + } + std::vector data; + Status download_status = dispatch_trace_download_binding(context, binding, data); + if (!download_status.success()) { + status.append(download_status); + continue; + } + const uint64_t hash = dispatch_trace_hash_bytes(data.data(), data.size()); + std::string blob_relative; + if (capture_blobs) { + std::ostringstream blob_name; + blob_name << stage << "-cmd" << command_index << "-ord" << command.ordinal << "-bind" << binding_index + << "-value" << binding.binding.value.value << "-off" << binding.binding.offset << "-" + << dispatch_trace_hex(hash) << ".bin"; + const std::filesystem::path blob_path = blobs / blob_name.str(); + status.append(dispatch_trace_write_blob(blob_path, data)); + blob_relative = std::filesystem::relative(blob_path, root).string(); + } + try { + std::filesystem::create_directories(root); + std::ofstream out(trace_path, std::ios::app); + if (!out) { + status.log("failed to open dispatch trace %s", trace_path.string().c_str()); + continue; + } + out << "{" + << "\"stage\":\"" << stage << "\"," + << "\"list_kind\":\"" << list_kind << "\"," + << "\"list_size\":" << list_size << ',' + << "\"program_command_count\":" << program_command_count << ',' + << "\"command_index\":" << command_index << ',' + << "\"ordinal\":" << command.ordinal << ',' + << "\"kernel_id\":" << command.kernel.specialization.kernel_id << ',' + << "\"kernel\":\"" << dispatch_trace_json_escape(kernel_name) << "\"," + << "\"binding_index\":" << binding_index << ',' + << "\"binding_name\":\"" << dispatch_trace_json_escape(binding.binding.name) << "\"," + << "\"value\":" << binding.binding.value.value << ',' + << "\"origin\":\"" << command_binding_origin_name(binding.binding.origin) << "\"," + << "\"access\":\"" << resource_access_name(binding.binding.access) << "\"," + << "\"offset\":" << binding.binding.offset << ',' + << "\"length\":" << binding.binding.length << ',' + << "\"ref_offset\":" << binding.ref.offset << ',' + << "\"ref_length\":" << binding.ref.length << ',' + << "\"hash\":\"" << dispatch_trace_hex(hash) << "\""; + if (!blob_relative.empty()) { + out << ",\"blob\":\"" << dispatch_trace_json_escape(blob_relative) << "\""; + } + out << "}\n"; + } catch (const std::exception & error) { + status.log("failed to write dispatch trace %s: %s", trace_path.string().c_str(), error.what()); + } + } + return status; +} + static Status command_program_metadata_context_valid(const CommandProgramExecutionContext & context) { Status status; if (context.target == nullptr) { @@ -1139,6 +1349,13 @@ static bool execute_prepared_command_list_serial(const CommandProgramExecutionCo if (!debug.sync("command-pre-dispatch", command_detail)) { return false; } + Status trace_status = + dispatch_trace_command_bindings(context, command, list_kind, commands.size(), i, program_command_count, + "before"); + if (!trace_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(trace_status)); + return false; + } if (trace_command) { debug.log("begin", "command-dispatch", command_detail); } @@ -1151,6 +1368,13 @@ static bool execute_prepared_command_list_serial(const CommandProgramExecutionCo if (!debug.sync("command-post-dispatch", command_detail)) { return false; } + trace_status = + dispatch_trace_command_bindings(context, command, list_kind, commands.size(), i, program_command_count, + "after"); + if (!trace_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(trace_status)); + return false; + } } if (trace_list) { debug.log("end", "command-list", @@ -1551,7 +1775,7 @@ static bool execute_prepared_command_prefix_via_graph(const CommandProgramExecut } // namespace bool debug_serial_command_execution_enabled() { - return environment_flag_enabled("GGML_HRX_DEBUG_SERIAL_EXECUTION"); + return environment_flag_enabled("GGML_HRX_DEBUG_SERIAL_EXECUTION") || dispatch_trace_enabled(); } PreparedCommandProgram prepare_command_program(const CommandProgramExecutionContext & context, diff --git a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py index a94bea2d966f..33753d93e748 100755 --- a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py +++ b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py @@ -715,11 +715,12 @@ def render_case(symbol: str, command: dict[str, Any], export: dict[str, Any]) -> elif kernel == "ggml_mul_mat_add_f32_f32_decode_wave64": case = case_mul_mat_add_decode(symbol, command, kernel) elif kernel in ( - "ggml_mul_mat_bias_f32_f32_wmma", - "ggml_mul_mat_add_f32_f32_wmma", - "ggml_mul_mat_bias_add_f32_f32_wmma", - "ggml_mul_mat_add_next_rmsnorm_f32_f32_wmma", - "ggml_mul_mat_bias_add_next_rmsnorm_f32_f32_wmma", + "ggml_mul_mat_tiled_input_f32_bias_publish_f32", + "ggml_mul_mat_tiled_input_f32_residual_publish_f32", + "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "ggml_mul_mat_skinny_input_f32_bias_publish_f32", + "ggml_mul_mat_skinny_input_f32_residual_publish_f32", + "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", ): case = case_mul_mat_postops(symbol, command, kernel) elif kernel in ("ggml_mul_mat_swiglu_f32_f32_wmma", "ggml_mul_mat_swiglu_f32_f32_decode_wave64"): diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index ebdc592c5dd2..e414ee9473dd 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -2648,9 +2648,12 @@ static void schedule_single_matmul_command(ggml_context * ctx, const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); const bool vector_kernel = string_contains(kernel_name, "ggml_mul_mat_vector"); const bool q6_vector_kernel = kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"; + const std::string expected_kernel(expected_kernel_name); + const bool q8_capable_kernel = expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || + expected_kernel == "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; const bool q8_input = !vector_kernel && expected_token_count <= 5 && expected_output_size % 64 == 0 && (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && - std::string(expected_kernel_name) == "loom_libs:ggml_mul_mat_f32_f32_wmma"; + q8_capable_kernel; const size_t command_count = q8_input ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); if (q8_input) { @@ -2658,7 +2661,6 @@ static void schedule_single_matmul_command(ggml_context * ctx, "qwen3_moe:ggml_quantize_q8_1_x4_f32"); } - const std::string expected_kernel(expected_kernel_name); const bool accepts_vector_token1 = expected_token_count == 1 && (expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" || @@ -5600,7 +5602,7 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { static void run_qwen_matmul_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 8 * 1024 * 1024; + params.mem_size = 12 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -5612,7 +5614,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5632,7 +5634,47 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 3); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 5); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 6); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 6, 2048, 128); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, 640, 128); } { @@ -5652,7 +5694,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5672,7 +5714,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5692,7 +5734,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5712,7 +5754,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5732,7 +5774,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -5752,7 +5794,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 256); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 256); } { @@ -6111,7 +6153,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_silu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", ggml::hrx::UnaryKind::Silu, 4, 2048, 128); } @@ -6124,7 +6166,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_relu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", ggml::hrx::UnaryKind::Relu, 4, 2048, 128); } @@ -6137,10 +6179,23 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_gelu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", ggml::hrx::UnaryKind::Gelu, 4, 2048, 128); } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_gelu(ctx, matmul); + REQUIRE(output != nullptr); + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + ggml::hrx::UnaryKind::Gelu, 2, 640, 128); + } + { ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); @@ -6182,7 +6237,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q4_input != nullptr); ggml_tensor * q4_output = ggml_mul_mat(ctx, q4_weight, q4_input); REQUIRE(q4_output != nullptr); - schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * q4_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); ggml_tensor * q4_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6201,7 +6256,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(legacy_input != nullptr); ggml_tensor * legacy_output = ggml_mul_mat(ctx, legacy_weight, legacy_input); REQUIRE(legacy_output != nullptr); - schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * legacy_decode_weight = ggml_new_tensor_2d(ctx, legacy_type, 2048, 128); ggml_tensor * legacy_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6219,7 +6274,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q3_input != nullptr); ggml_tensor * q3_output = ggml_mul_mat(ctx, q3_weight, q3_input); REQUIRE(q3_output != nullptr); - schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * q3_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q3_K, 2048, 128); ggml_tensor * q3_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6236,7 +6291,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_input != nullptr); ggml_tensor * iq2_s_output = ggml_mul_mat(ctx, iq2_s_weight, iq2_s_input); REQUIRE(iq2_s_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 640); + schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 640); ggml_tensor * iq2_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ2_S, 2048, 640); ggml_tensor * iq2_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6253,7 +6308,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_input != nullptr); ggml_tensor * iq3_s_output = ggml_mul_mat(ctx, iq3_s_weight, iq3_s_input); REQUIRE(iq3_s_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * iq3_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); ggml_tensor * iq3_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6270,7 +6325,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_nl_input != nullptr); ggml_tensor * iq4_nl_output = ggml_mul_mat(ctx, iq4_nl_weight, iq4_nl_input); REQUIRE(iq4_nl_output != nullptr); - schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * iq4_nl_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); ggml_tensor * iq4_nl_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6287,7 +6342,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_nl_small_input != nullptr); ggml_tensor * iq4_nl_small_output = ggml_mul_mat(ctx, iq4_nl_small_weight, iq4_nl_small_input); REQUIRE(iq4_nl_small_output != nullptr); - schedule_single_matmul_command(ctx, iq4_nl_small_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640, 1024); + schedule_single_matmul_command(ctx, iq4_nl_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + 2, 640, 1024); ggml_tensor * iq4_nl_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 1024); ggml_tensor * iq4_nl_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); @@ -6305,7 +6361,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q5_0_small_input != nullptr); ggml_tensor * q5_0_small_output = ggml_mul_mat(ctx, q5_0_small_weight, q5_0_small_input); REQUIRE(q5_0_small_output != nullptr); - schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640, 256); + schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + 2, 640, 256); ggml_tensor * q5_0_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); ggml_tensor * q5_0_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); @@ -6322,7 +6379,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q6_input != nullptr); ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, q6_input); REQUIRE(q6_output != nullptr); - schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 128); + schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, 128); ggml_tensor * q6_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); ggml_tensor * q6_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6348,7 +6405,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_xs_input != nullptr); ggml_tensor * iq4_xs_output = ggml_mul_mat(ctx, iq4_xs_weight, iq4_xs_input); REQUIRE(iq4_xs_output != nullptr); - schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * iq4_xs_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); ggml_tensor * iq4_xs_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6364,7 +6421,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_0_input != nullptr); ggml_tensor * q8_0_output = ggml_mul_mat(ctx, q8_0_weight, q8_0_input); REQUIRE(q8_0_output != nullptr); - schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * q8_0_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); ggml_tensor * q8_0_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6399,7 +6456,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_1_input != nullptr); ggml_tensor * q8_1_output = ggml_mul_mat(ctx, q8_1_weight, q8_1_input); REQUIRE(q8_1_output != nullptr); - schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * q8_1_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); ggml_tensor * q8_1_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6416,7 +6473,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f16_input != nullptr); ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, f16_input); REQUIRE(f16_output != nullptr); - schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * f16_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); ggml_tensor * f16_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6433,7 +6490,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f32_input != nullptr); ggml_tensor * f32_output = ggml_mul_mat(ctx, f32_weight, f32_input); REQUIRE(f32_output != nullptr); - schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); ggml_tensor * f32_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); ggml_tensor * f32_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -6487,7 +6544,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, 128); } { @@ -10482,16 +10539,18 @@ static void run_lowtoken_residual_dispatch_checks() { for (const auto & command : program.commands) { const std::string name = kernel_name_for_id(command.kernel.kernel_id); binary += name == "loom_libs:ggml_binary_f32"; - const bool legacy_postops = name == "loom_libs:ggml_mul_mat_bias_f32_f32_wmma" || - name == "loom_libs:ggml_mul_mat_add_f32_f32_wmma"; const bool tiled_postops = name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32" || name == "loom_libs:ggml_mul_mat_tiled_input_f32_residual_publish_f32" || name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; + const bool skinny_postops = + name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_publish_f32" || + name == "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32" || + name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_f32_f32" || name == "loom_libs:ggml_mul_mat_vector_residual_f32_f32" || name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; - if (!legacy_postops && !tiled_postops && !vector_postops) { + if (!tiled_postops && !skinny_postops && !vector_postops) { continue; } ++fused; @@ -10512,6 +10571,41 @@ static void run_lowtoken_residual_dispatch_checks() { } } +static void run_lowtoken_bias_dispatch_checks() { + ggml_context * ctx = ggml_init({ 16 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 3584, 1024); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3584, 2); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1024); + ggml_tensor * output = ggml_add(ctx, projection, bias); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandProgram program = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(program.valid()); + REQUIRE(command_program_verifies(program)); + + size_t skinny_bias = 0; + size_t tiled_bias = 0; + for (const auto & command : program.commands) { + const std::string name = kernel_name_for_id(command.kernel.kernel_id); + skinny_bias += name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_publish_f32"; + tiled_bias += name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_publish_f32"; + } + REQUIRE(skinny_bias == 1); + REQUIRE(tiled_bias == 0); + + ggml_free(ctx); +} + static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 32 * 1024 * 1024; @@ -10704,6 +10798,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("rmsnorm_gate_packed_output", [] { run_rmsnorm_gate_packed_output_checks(); }); suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); suite.host_case("lowtoken_residual_dispatch", [] { run_lowtoken_residual_dispatch_checks(); }); + suite.host_case("lowtoken_bias_dispatch", [] { run_lowtoken_bias_dispatch_checks(); }); suite.host_case("quantized_conv4_dispatch", [] { run_quantized_conv4_dispatch_checks(); }); suite.host_case("gdn_rmsnorm_gate_dispatch", [] { run_gdn_rmsnorm_gate_dispatch_checks(); }); suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index c47f46a2ecd2..6b09f5af0f82 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -4680,7 +4680,8 @@ static void register_dense_matmul_cases(Suite & suite) { [tokens, outputs] { run_dense_matmul_cpu_reference_case( GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, outputs); }); } @@ -4688,16 +4689,18 @@ static void register_dense_matmul_cases(Suite & suite) { for (const ggml_type type : { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { suite.device_case("dense_matmul." + type_name(type) + ".tokens3.outputs47", [type] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 47); }); } suite.device_case("dense_matmul_unary.f32.tokens3.outputs47", [] { - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + 3, 47); }); suite.device_case("get_rows_f32.f32.rows262144.tokens4.outputs1", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); }); - suite.device_case("dense_matmul.q4_k.wmma.tokens2.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + suite.device_case("dense_matmul.q4_k.skinny.tokens2.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, + 128); }); suite.device_case("dense_matmul.q4_k.vector.tokens1.outputs128", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); @@ -4719,38 +4722,46 @@ static void register_dense_matmul_cases(Suite & suite) { "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; - suite.device_case("dense_matmul." + type_name(type) + ".wmma.tokens2.outputs" + std::to_string(outputs), - [type, outputs, input] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, - outputs, input); + const char * route = type == GGML_TYPE_IQ4_NL ? "tiled" : "skinny"; + const char * kernel = 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); }); suite.device_case("dense_matmul." + type_name(type) + ".vector.tokens1.outputs" + std::to_string(outputs), [type, vector_kernel, outputs, input] { run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); }); } - suite.device_case("dense_matmul.q5_0.wmma.tokens2.outputs256.input640", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256, 640); + suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs256.input640", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, + 256, 640); }); - suite.device_case("dense_matmul.q5_0.wmma.tokens2.outputs128.input544", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128, 544); + suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs128.input544", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, + 128, 544); }); suite.device_case("dense_matmul.q5_0.vector.tokens1.outputs256.input640", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); }); - suite.device_case("dense_matmul.q8_1.zero_weight.wmma", - [] { run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_wmma"); }); + suite.device_case("dense_matmul.q8_1.zero_weight.skinny", [] { + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"); + }); suite.device_case("dense_matmul.q8_1.zero_weight.vector", [] { run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); }); - suite.device_case("dense_matmul.f32.wmma.tokens2.outputs256", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); + suite.device_case("dense_matmul.f32.skinny.tokens2.outputs256", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, + 256); }); suite.device_case("dense_matmul.f32.vector.tokens1.outputs256", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); }); suite.device_case("dense_matmul_unary.f32.tokens2.outputs128", [] { - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + 2, 128); }); suite.device_case("dense_matmul_swiglu.q4_k.f16.swiglu.tokens2.outputs128", [] { run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, @@ -4883,72 +4894,86 @@ static void register_dense_matmul_cases(Suite & suite) { for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; suite.device_case( - "dense_matmul." + type_name(type) + ".wmma.tokens5.outputs4096.input" + std::to_string(input_size), + "dense_matmul." + type_name(type) + ".skinny.tokens5.outputs4096.input" + std::to_string(input_size), [type, input_size] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, input_size); + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, input_size); }); suite.device_case("dense_matmul_postops." + type_name(type) + ".add.tokens5.outputs4160", [type, input_size] { - run_dense_matmul_postops_cpu_reference_case(type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, - true, false, input_size + 256); + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4160, false, true, false, + input_size + 256); }); suite.device_case( "dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", [type, input_size] { - run_dense_matmul_postops_cpu_reference_case(type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, - false, true, true, input_size + 256); + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4224, false, true, true, + input_size + 256); }); } - suite.device_case("dense_matmul.q4_k.wmma.tokens5.outputs4096.input8192", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 8192); + suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs4096.input8192", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, 8192); }); for (const int64_t input_size : { 6144, 6400, 6656 }) { - suite.device_case("dense_matmul.q4_k.wmma.tokens5.outputs5120.input" + std::to_string(input_size), + suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", - 5, 5120, input_size); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 5120, + input_size); }); suite.device_case( "dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", - 5, 5120, false, true, true, input_size); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 5120, false, + true, true, input_size); }); } - suite.device_case("dense_matmul.q4_k.wmma.tokens4.outputs4096.input4096", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 4096, 4096); + suite.device_case("dense_matmul.q4_k.skinny.tokens4.outputs4096.input4096", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 4096, 4096); }); suite.device_case("dense_matmul_postops.q4_k.add.tokens4.outputs4160.input6400", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4160, - false, true, false, 6400); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 4, 4160, false, true, + false, 6400); }); suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens4.outputs4224.input6656", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4224, - false, true, true, 6656); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 4, 4224, false, true, + true, 6656); }); suite.device_case("endpoint_rmsnorm_q6k_q8.cpu_reference", [] { run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); }); - suite.device_case("dense_matmul.q6_k.wmma.tokens5.outputs4096.input16384", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 16384); + suite.device_case("dense_matmul.q6_k.skinny.tokens5.outputs4096.input16384", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, 16384); }); suite.device_case("dense_matmul_postops.q6_k.add.tokens5.outputs4160.input16640", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, - false, true, false, 16640); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4160, false, true, + false, 16640); }); suite.device_case("dense_matmul_postops.q6_k.add.alias.tokens5.outputs4224.input16640", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, - false, true, true, 16640); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_residual_publish_f32", 5, 4224, false, true, + true, 16640); }); for (const int64_t tokens : { 1, 3, 5 }) { suite.device_case( "dense_matmul.q4_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { run_dense_matmul_cpu_reference_case( GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, 256, 5120, true); }); suite.device_case( "dense_matmul.q6_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { run_dense_matmul_cpu_reference_case( GGML_TYPE_Q6_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", tokens, 256, 5120, true, true); }); }