diff --git a/ggml/src/ggml-hrx/CMakeLists.txt b/ggml/src/ggml-hrx/CMakeLists.txt index 805545884c96..af9dde2263b3 100644 --- a/ggml/src/ggml-hrx/CMakeLists.txt +++ b/ggml/src/ggml-hrx/CMakeLists.txt @@ -220,6 +220,8 @@ ggml_add_backend_library(ggml-hrx dispatch_registration/common/dispatch-get-rows.h dispatch_registration/common/dispatch-glu.cpp dispatch_registration/common/dispatch-glu.h + dispatch_registration/common/dispatch-hyper-connection.cpp + dispatch_registration/common/dispatch-hyper-connection.h dispatch_registration/common/dispatch-mul-mat-id.cpp dispatch_registration/common/dispatch-mul-mat-id-common.h dispatch_registration/common/dispatch-mul-mat-id.h @@ -235,6 +237,8 @@ ggml_add_backend_library(ggml-hrx dispatch_registration/common/dispatch-scale.cpp dispatch_registration/common/dispatch-scale.h dispatch_registration/common/dispatch-unary.cpp + dispatch_registration/common/dispatch-reduction.cpp + dispatch_registration/common/dispatch-reduction.h dispatch_registration/common/dispatch-unary.h dispatch_registration/common/iq-codebook-resource.cpp dispatch_registration/common/iq-codebook-resource.h diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md index e3152cfaf117..e6897fa21a66 100644 --- a/ggml/src/ggml-hrx/README.md +++ b/ggml/src/ggml-hrx/README.md @@ -52,6 +52,8 @@ jq -c 'select(.event == "compile_request") | {cache_id, execution_id, kernel, in `hrx-minimum-commit.txt` records the minimum required commit in ROCm/hrx-system. It contains one full Git commit SHA. The backend's C++ APIs, Loom syntax, and numerical behavior can all depend on that revision. When `GGML_HRX=ON` and `HRX_SOURCE_DIR` is set, CMake requires the HRX checkout's HEAD to equal or descend from the recorded commit. +The corpus uses typed `check.expect.approximate` operations introduced in [HRX d03575e84](https://github.com/ROCm/hrx-system/commit/d03575e84f602d3e39b87895a30fae214f308558). + For builds using installed packages, check the source revision used to produce both HRX and Loom against this floor before validating a repair. From the staging checkout: ```sh diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp index 12d4f5fa338a..cc607daf8adb 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp @@ -135,7 +135,8 @@ static bool can_elide_zero_output_node(const Graph & graph, } static bool can_elide_node(const Graph & graph, const GraphNode & node, const std::vector & covered_nodes) { - return can_elide_layout_alias_node(graph, node, covered_nodes) || + return (node.op == GGML_OP_NONE && node.inputs.empty()) || + can_elide_layout_alias_node(graph, node, covered_nodes) || can_elide_zero_output_node(graph, node, covered_nodes); } @@ -204,10 +205,10 @@ bool DispatchScheduler::schedule_graph(Graph & graph, covered_nodes[i] = true; continue; } - plan_.status.append(pending_diagnostics); - plan_.status.append(match.status); const std::string message = unsupported_node_message(graph, i, *node); plan_.status.log("%s", message.c_str()); + plan_.status.append(match.status); + plan_.status.append(pending_diagnostics); if (diagnostics != nullptr) { diagnostics->unsupported_node_index = i; diagnostics->unsupported_node = node; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp index c42ac59a441a..5c14fb089ce0 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp @@ -193,18 +193,14 @@ static void add_binary_strided_parameters(Dispatch & dispatch, dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(output.ne[0])); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(output.ne[1])); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", std::to_string(output.ne[2])); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", - std::to_string(lhs.nb[1] / sizeof(float))); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", - std::to_string(lhs.nb[2] / sizeof(float))); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", - std::to_string(lhs.nb[3] / sizeof(float))); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", - std::to_string(rhs.nb[1] / sizeof(float))); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", - std::to_string(rhs.nb[2] / sizeof(float))); - dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", - std::to_string(rhs.nb[3] / sizeof(float))); + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + dispatch.kernel.compile_parameters.emplace( + "ggml.binary_f32.src0_stride" + std::to_string(dim), + std::to_string(broadcast_dim_flag(lhs, output, dim) ? 0 : lhs.nb[dim] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace( + "ggml.binary_f32.src1_stride" + std::to_string(dim), + std::to_string(broadcast_dim_flag(rhs, output, dim) ? 0 : rhs.nb[dim] / sizeof(float))); + } dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", std::to_string(lhs_byte_count / sizeof(float))); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", @@ -298,9 +294,14 @@ static bool match_binary_f32_dispatch(const DispatchMatchContext & context, Disp } const bool same_shape_inputs = same_shape(*lhs, *output) && same_shape(*rhs, *output); + const bool strided_broadcast = !same_shape_inputs && (!lhs->contiguous || !rhs->contiguous || + !packed_f32_layout(*lhs) || !packed_f32_layout(*rhs)); + if (!binary_kind_allows_broadcast(params->op, *lhs, *rhs, *output)) { + return false; + } const CommonActivationPublicationCapabilities capabilities = { - binary_q8_publication_enabled(params->op, *output) ? DispatchActivationInputQ8_1X4 : - DispatchActivationInputNone, + !strided_broadcast && binary_q8_publication_enabled(params->op, *output) ? DispatchActivationInputQ8_1X4 : + DispatchActivationInputNone, DispatchActivationInputQ8_1X4, }; const CommonActivationPublicationPlan publication_plan = @@ -310,22 +311,21 @@ static bool match_binary_f32_dispatch(const DispatchMatchContext & context, Disp const bool publish_q8 = q8_demand != nullptr; Dispatch dispatch; - if (same_shape_inputs) { + if (same_shape_inputs || strided_broadcast) { dispatch.kernel = make_kernel_specialization(publish_q8 ? kBinaryF32PublishQ8_1X4Kernel : kBinaryF32Kernel); add_binary_strided_parameters(dispatch, *lhs, *rhs, *output, lhs_byte_count, rhs_byte_count); if (publish_q8) { dispatch.kernel.integer_parameters.erase("element_count"); dispatch.kernel.integer_parameters.emplace("token_count", output->ne[1]); + } else { + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride0", + std::to_string(broadcast_dim_flag(*lhs, *output, 0) ? 0 : 1)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride0", + std::to_string(broadcast_dim_flag(*rhs, *output, 0) ? 0 : 1)); } dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.op", std::to_string(binary_kind_config_value(params->op))); } else { - if (!lhs->contiguous || !rhs->contiguous || !packed_f32_layout(*lhs) || !packed_f32_layout(*rhs)) { - return false; - } - if (!binary_kind_allows_broadcast(params->op, *lhs, *rhs, *output)) { - return false; - } dispatch.kernel = make_kernel_specialization(publish_q8 ? kBinaryBcF32PublishQ8_1X4Kernel : kBinaryBcF32Kernel); if (publish_q8) { dispatch.kernel.integer_parameters.emplace("token_count", output->ne[1]); @@ -356,6 +356,33 @@ static bool match_binary_f32_dispatch(const DispatchMatchContext & context, Disp return match.status.success(); } +static bool match_add_f16(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_ADD || node->inputs.size() != 2) { + return false; + } + const Value * lhs = graph_value(context.graph, node->inputs[0]); + const Value * rhs = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (lhs == nullptr || rhs == nullptr || output == nullptr || lhs->type != GGML_TYPE_F16 || + rhs->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F16 || !lhs->contiguous || !rhs->contiguous || + !output->contiguous || lhs->ne != output->ne || rhs->ne != output->ne || output->element_count < 1 || + output->element_count > 134217728 || !binary_output_storage_is_safe(*lhs, *rhs, *output)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_add_f16")); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.bindings = { + { lhs->id, 0, lhs->byte_count }, + { rhs->id, 0, rhs->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + static void register_binary_dispatch_for(DispatchRegistryBuilder & registry, ggml_op root_op) { registry.add({ "common.binary_f32", @@ -370,6 +397,8 @@ static void register_binary_dispatch_for(DispatchRegistryBuilder & registry, ggm } // namespace void register_binary_dispatch(DispatchRegistryBuilder & registry) { + registry.add( + { "common.add.f16", GGML_OP_ADD, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, match_add_f16 }); registry.add({ "common.binary_swiglu_symmetric_i4_k32", GGML_OP_GLU, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp index 9ccaebe22270..60e10fc1d303 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp @@ -8,15 +8,27 @@ #include "dispatch-gather-add.h" #include "dispatch-get-rows.h" #include "dispatch-glu.h" +#include "dispatch-hyper-connection.h" #include "dispatch-mul-mat-id.h" #include "dispatch-mul-mat.h" +#include "dispatch-reduction.h" #include "dispatch-rmsnorm.h" #include "dispatch-rope-set-rows.h" #include "dispatch-scale.h" #include "dispatch-unary.h" +#include +#include + namespace ggml::hrx { +std::string float_config_value(float value) { + std::ostringstream out; + out.precision(std::numeric_limits::max_digits10); + out << value; + return out.str(); +} + void register_common_dispatches(DispatchRegistryBuilder & registry) { register_binary_dispatch(registry); register_copy_dispatch(registry); @@ -26,8 +38,10 @@ void register_common_dispatches(DispatchRegistryBuilder & registry) { register_gather_add_dispatch(registry); register_get_rows_dispatches(registry); register_glu_dispatches(registry); + register_hyper_connection_dispatches(registry); register_mul_mat_id_dispatches(registry); register_mul_mat_dispatches(registry); + register_reduction_dispatches(registry); register_rope_set_rows_dispatches(registry); register_scale_dispatch(registry); register_unary_dispatch(registry); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h index d7fe2aed8246..88a5bfdee590 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h @@ -2,8 +2,12 @@ #include "../dispatch-registry.h" +#include + namespace ggml::hrx { +std::string float_config_value(float value); + void register_common_dispatches(DispatchRegistryBuilder & registry); } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp index 37c9f1ffe893..17087e7424f2 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp @@ -61,7 +61,9 @@ static void add_concat_strided_source_parameters(KernelSpecialization & kernel, } } -static bool make_copy_f32_dispatch(const Value & source, const Value & output, Dispatch & dispatch) { +} // namespace + +bool make_copy_f32_dispatch(const Value & source, const Value & output, Dispatch & dispatch) { if (source.type != GGML_TYPE_F32 || output.type != GGML_TYPE_F32 || !output.contiguous || !packed_f32_layout(output) || !f32_element_strides(source) || source.element_count <= 0 || source.element_count != output.element_count || @@ -74,7 +76,7 @@ static bool make_copy_f32_dispatch(const Value & source, const Value & output, D dispatch.kernel = make_kernel_specialization(kCopyF32Kernel); dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; } else { - if (!strided_f32_storage_span_bytes(source, source_span)) { + if (!strided_f32_storage_span_bytes(source, source_span, true)) { return false; } const size_t source_span_elements = source_span / sizeof(float); @@ -97,6 +99,69 @@ static bool make_copy_f32_dispatch(const Value & source, const Value & output, D return true; } +namespace { + +static bool make_copy_bits_dispatch(const Value & source, + const Value & output, + Dispatch & dispatch, + bool repeat = false) { + if (source.type != output.type || + (source.type != GGML_TYPE_F16 && source.type != GGML_TYPE_I32 && source.type != GGML_TYPE_F32) || + !output.contiguous || source.element_count < 1 || (!repeat && source.element_count != output.element_count) || + output.element_count < 1 || output.element_count > static_cast(kMaximumCopyElements) || + source.storage == output.storage) { + return false; + } + const size_t element_size = ggml_type_size(source.type); + if (source.byte_count / element_size > kMaximumCopyElements) { + return false; + } + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (source.ne[dim] < 1 || source.nb[dim] == 0 || source.nb[dim] % element_size != 0 || + source.nb[dim] / element_size > kMaximumCopyElements) { + return false; + } + } + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_bits")); + dispatch.kernel.integer_parameters.emplace("element_count", output.element_count); + dispatch.kernel.integer_parameters.emplace("source_span", source.byte_count / element_size); + auto & config = dispatch.kernel.compile_parameters; + config.emplace("ggml.copy_bits.element_bits", std::to_string(element_size * 8)); + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (repeat && (output.ne[dim] < 1 || output.ne[dim] % source.ne[dim] != 0)) { + return false; + } + if (dim < 3) { + config.emplace("ggml.copy_bits.ne" + std::to_string(dim), + std::to_string(repeat ? output.ne[dim] : source.ne[dim])); + } + config.emplace("ggml.copy_bits.source_ne" + std::to_string(dim), std::to_string(source.ne[dim])); + config.emplace("ggml.copy_bits.stride" + std::to_string(dim), std::to_string(source.nb[dim] / element_size)); + } + dispatch.bindings = { + { source.id, 0, source.byte_count }, + { output.id, 0, output.byte_count } + }; + return true; +} + +static bool match_cont_bits_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || (node->op != GGML_OP_CONT && node->op != GGML_OP_REPEAT) || node->inputs.size() != 1) { + return false; + } + const Value * source = graph_value(context.graph, node->inputs[0]); + const Value * output = graph_value(context.graph, node->output); + Dispatch dispatch; + if (source == nullptr || output == nullptr || + !make_copy_bits_dispatch(*source, *output, dispatch, node->op == GGML_OP_REPEAT)) { + return false; + } + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + static bool match_copy_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { const GraphNode * node = context.root_node; if (node == nullptr || node->op != GGML_OP_CPY || node->inputs.size() != 2) { @@ -209,7 +274,69 @@ static bool match_concat_dim0_f32_dispatch(const DispatchMatchContext & context, } // namespace +static bool match_concat_packed(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2) { + return false; + } + const Value * lhs = graph_value(context.graph, node->inputs[0]); + const Value * rhs = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (lhs == nullptr || rhs == nullptr || output == nullptr || + (lhs->type != GGML_TYPE_F16 && lhs->type != GGML_TYPE_F32 && lhs->type != GGML_TYPE_I32) || + rhs->type != lhs->type || output->type != lhs->type || !lhs->contiguous || !rhs->contiguous || + !output->contiguous || output->element_count < 1 || output->element_count > 134217728 || + lhs->storage == output->storage || rhs->storage == output->storage) { + return false; + } + int axis = -1; + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (lhs->ne[dim] < 1 || rhs->ne[dim] < 1) { + return false; + } + if (output->ne[dim] == lhs->ne[dim] + rhs->ne[dim]) { + if (axis >= 0) { + return false; + } + axis = dim; + } else if (lhs->ne[dim] != output->ne[dim] || rhs->ne[dim] != output->ne[dim]) { + return false; + } + } + if (axis < 0 || (axis == 0 && lhs->type == GGML_TYPE_F32)) { + return false; + } + int64_t lhs_chunk = 1, rhs_chunk = 1; + for (int dim = 0; dim <= axis; ++dim) { + lhs_chunk *= lhs->ne[dim]; + rhs_chunk *= rhs->ne[dim]; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_concat_packed")); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.concat_packed.element_bits", + std::to_string(8 * ggml_type_size(lhs->type))); + dispatch.kernel.compile_parameters.emplace("ggml.concat_packed.lhs_chunk", std::to_string(lhs_chunk)); + dispatch.kernel.compile_parameters.emplace("ggml.concat_packed.rhs_chunk", std::to_string(rhs_chunk)); + dispatch.kernel.compile_parameters.emplace("ggml.concat_packed.lhs_span", std::to_string(lhs->element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.concat_packed.rhs_span", std::to_string(rhs->element_count)); + dispatch.bindings = { + { lhs->id, 0, lhs->byte_count }, + { rhs->id, 0, rhs->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + void register_copy_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ "common.repeat_bits", GGML_OP_REPEAT, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_cont_bits_dispatch }); + registry.add({ "common.concat_packed", GGML_OP_CONCAT, DispatchMatchKind::SingleOp, 10, DispatchSource::Common, + match_concat_packed }); + registry.add({ "common.cont_bits", GGML_OP_CONT, DispatchMatchKind::SingleOp, -10, DispatchSource::Common, + match_cont_bits_dispatch }); registry.add({ "common.copy_f32", GGML_OP_CPY, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h index 6228aac1946b..33c7a8abfaff 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h @@ -4,6 +4,8 @@ namespace ggml::hrx { +bool make_copy_f32_dispatch(const Value & source, const Value & output, Dispatch & dispatch); + void register_copy_dispatch(DispatchRegistryBuilder & registry); } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp index c1fec1a8a8ce..79b70e6586bd 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp @@ -17,12 +17,16 @@ namespace { static constexpr KernelCatalogRef kFlashAttentionF32F16WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_f32_f16_wmma"); +static constexpr KernelCatalogRef kFlashAttentionF32F16WmmaSinkKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_f32_f16_wmma_sink"); static constexpr KernelCatalogRef kFlashAttentionF32F16WmmaPublishF16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_f32_f16_wmma_publish_f16"); static constexpr KernelCatalogRef kFlashAttentionDecodeSplitNextQ8Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); static constexpr KernelCatalogRef kCopyTransposeF16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_transpose_f16"); +static constexpr KernelCatalogRef kLightningIndexerKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_lightning_indexer_f32_f16"); static constexpr int64_t kDecodeRowCapacity = 16; static constexpr int64_t kDecodeKvTileSize = 64; static constexpr int64_t kPrefillQkHeadSizeBlock = 16; @@ -36,13 +40,60 @@ static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); } +static bool match_lightning_indexer(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_LIGHTNING_INDEXER || node->inputs.size() != 4) { + return false; + } + const Value * q = graph_value(context.graph, node->inputs[0]); + const Value * k = graph_value(context.graph, node->inputs[1]); + const Value * w = graph_value(context.graph, node->inputs[2]); + const Value * m = graph_value(context.graph, node->inputs[3]); + const Value * y = graph_value(context.graph, node->output); + if (!q || !k || !w || !m || !y || !q->contiguous || !k->contiguous || !w->contiguous || !m->contiguous || + !y->contiguous || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || w->type != GGML_TYPE_F32 || + m->type != GGML_TYPE_F16 || y->type != GGML_TYPE_F32) { + return false; + } + const int64_t nq = q->ne[2]; + const int64_t nk = k->ne[2]; + const int64_t head_size = q->ne[0]; + const int64_t head_count = q->ne[1]; + // Keep staged operands within the existing workgroup memory budget. + if (nq < 1 || nq > 2048 || nk < 1 || nk > 131072 || head_size < 16 || head_size > 256 || head_size % 16 != 0 || + head_count < 1 || head_count > 64 || ((head_count + 15) / 16) * 16 * head_size > 8192 || q->ne[3] != 1 || + k->ne != std::array{ head_size, 1, nk, 1 } || + w->ne != std::array{ head_count, nq, 1, 1 } || m->ne != std::array{ nk, nq, 1, 1 } || + y->ne != std::array{ nk, nq, 1, 1 }) { + return false; + } + for (const Value * input : { q, k, w, m }) { + if (input->storage == y->storage) { + return false; + } + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kLightningIndexerKernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + dispatch.kernel.integer_parameters.emplace("query_count", nq); + dispatch.kernel.integer_parameters.emplace("key_count", nk); + dispatch.kernel.compile_parameters.emplace("ggml.indexer.head_size", std::to_string(head_size)); + dispatch.kernel.compile_parameters.emplace("ggml.indexer.head_count", std::to_string(head_count)); + for (const Value * value : { q, k, w, m, y }) { + dispatch.bindings.push_back({ value->id, 0, value->byte_count }); + } + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool nearly_equal(float lhs, float rhs) { return std::fabs(lhs - rhs) <= 1.0e-6f; } static bool use_exact_prefill_workload(int64_t query_token_count, int64_t key_value_token_count) { - // Preserve the measured hot prefill case; other supported totals share dynamic executables. - return query_token_count == 256 && key_value_token_count == 256; + // Aligned prefill has no partial KV tile; specialize away its scratch-using cleanup. + return query_token_count > 1 && key_value_token_count % kDecodeKvTileSize == 0; } static bool is_supported_token_count(int64_t token_count) { @@ -143,6 +194,7 @@ struct FlashAttentionMatch { const Value * value = nullptr; const Value * mask = nullptr; const Value * output = nullptr; + const Value * sink = nullptr; const GraphNode * output_layout = nullptr; ValueId mask_binding_value; size_t mask_binding_bytes = 0; @@ -182,9 +234,11 @@ struct DecodeSplitFlashAttentionMatch { static FlashAttentionMatch match_flash_attention_f32_f16(const Graph & graph, const CommandPlan & plan, - const GraphNode * node) { + const GraphNode * node, + bool allow_sparse_long_context = false) { FlashAttentionMatch match; - if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4) { + if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || + (node->inputs.size() != 4 && node->inputs.size() != 5)) { return match; } @@ -216,9 +270,16 @@ static FlashAttentionMatch match_flash_attention_f32_f16(const Graph & gra const int64_t query_token_count = query->ne[1]; const int64_t query_head_count = query->ne[2]; + if (node->inputs.size() == 5) { + match.sink = graph_value(graph, node->inputs[4]); + if (match.sink == nullptr || match.sink->type != GGML_TYPE_F32 || !match.sink->contiguous || + match.sink->element_count != query_head_count) { + return {}; + } + } const int64_t key_value_capacity = key->ne[1]; const int64_t key_value_head_count = key->ne[2]; - if (!is_supported_token_count(query_token_count) || key_value_capacity < query_token_count || + if (!is_supported_token_count(query_token_count) || key_value_capacity < 1 || !is_supported_head_count(query_head_count) || !is_supported_head_count(key_value_head_count) || query_head_count % key_value_head_count != 0) { return {}; @@ -245,7 +306,9 @@ static FlashAttentionMatch match_flash_attention_f32_f16(const Graph & gra } else if (!mask_is_compact && !mask_is_capacity) { return {}; } - if (!is_supported_key_value_token_count(key_value_token_count)) { + const bool sparse_long_context = allow_sparse_long_context && key_value_head_count == 1 && + key_value_token_count >= 1 && key_value_token_count <= 131072; + if (!sparse_long_context && !is_supported_key_value_token_count(key_value_token_count)) { return {}; } @@ -426,7 +489,8 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext static bool match_flash_attention_gate_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const FlashAttentionMatch match = match_flash_attention_f32_f16(context.graph, context.plan, context.root_node); - if (!match.matched() || match.output_layout == nullptr || match.output_layout->op != GGML_OP_RESHAPE) { + if (!match.matched() || match.sink != nullptr || match.output_layout == nullptr || + match.output_layout->op != GGML_OP_RESHAPE) { return false; } @@ -510,6 +574,88 @@ static bool match_flash_attention_gate_dispatch(const DispatchMatchContext & con return true; } +static bool match_sparse_attention(const DispatchMatchContext & context, DispatchMatch & result) { + const FlashAttentionMatch match = + match_flash_attention_f32_f16(context.graph, context.plan, context.root_node, true); + if (!match.matched() || match.sink == nullptr || match.query_head_count != 64 || match.key_value_head_count != 1 || + match.qk_head_size != 512 || match.value_head_size != 512 || match.key->storage != match.value->storage || + match.key->storage_offset != match.value->storage_offset || match.key->nb != match.value->nb) { + return false; + } + const bool prefill = match.query_token_count > 1; + const auto * params = op_params_as(context.root_node->params); + if (params == nullptr || params->n_kv_max <= 0 || params->n_kv_max > 2048) { + return false; + } + const int64_t capacity = ceil_div(params->n_kv_max, 64) * 64; + if (match.key_value_token_count < 2048 || match.key_value_token_count < 2 * capacity) { + return false; + } + const size_t index_bytes = static_cast(match.query_token_count * capacity) * sizeof(int32_t); + const size_t mask_bytes = static_cast(match.query_token_count * capacity) * sizeof(ggml_fp16_t); + const size_t key_bytes = mask_bytes * match.qk_head_size; + const ValueId indices = match_value(context, result, 0); + const ValueId mask = match_value(context, result, 1); + const ValueId key = match_value(context, result, 2); + result.transients.push_back({ indices, "common.sparse_attention.indices", index_bytes, 256 }); + result.transients.push_back({ mask, "common.sparse_attention.mask", mask_bytes, 256 }); + result.transients.push_back({ key, "common.sparse_attention.key", key_bytes, 256 }); + + Dispatch compact; + compact.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_compact_mask_f16")); + compact.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + compact.kernel.integer_parameters.emplace("query_count", match.query_token_count); + compact.kernel.integer_parameters.emplace("key_count", match.key_value_token_count); + compact.kernel.compile_parameters.emplace("ggml.sparse_attention.capacity", std::to_string(capacity)); + compact.bindings = { + { match.mask_binding_value, 0, match.mask_binding_bytes }, + { indices, 0, index_bytes }, + { mask, 0, mask_bytes } + }; + result.dispatches.push_back(std::move(compact)); + + Dispatch gather; + gather.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_gather_sparse_kv_f16")); + gather.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + gather.kernel.integer_parameters.emplace("query_count", match.query_token_count); + gather.kernel.integer_parameters.emplace("key_count", match.key_value_token_count); + gather.kernel.compile_parameters.emplace("ggml.sparse_attention.capacity", std::to_string(capacity)); + gather.bindings = { + { indices, 0, index_bytes }, + { match.key->id, 0, match.key->byte_count }, + { key, 0, key_bytes } + }; + result.dispatches.push_back(std::move(gather)); + + Dispatch attention; + attention.kernel = + make_kernel_specialization(prefill ? GGML_HRX_KERNEL_REF("loom_libs", "ggml_sparse_prefill_f32_f16") : + kFlashAttentionF32F16WmmaSinkKernel); + if (prefill) { + attention.kernel.integer_parameters.emplace("query_count", match.query_token_count); + attention.kernel.compile_parameters.emplace("ggml.sparse_attention.capacity", std::to_string(capacity)); + } else { + attention.kernel.integer_parameters.emplace("query_token_count", 4); + attention.kernel.integer_parameters.emplace("key_value_token_count", capacity); + } + add_flash_attention_compile_parameters(attention.kernel, prefill ? 4 : 16, 1, 512, 512, match.attention_scale, + false, 1, 1); + attention.kernel.compile_parameters.emplace("ggml.flash_attention.head_batch_size", prefill ? "16" : "4"); + attention.bindings = { + { match.query->id, 0, match.query->byte_count }, + { key, 0, key_bytes } + }; + if (!prefill) { + attention.bindings.push_back({ key, 0, key_bytes }); + } + attention.bindings.push_back({ mask, 0, mask_bytes }); + attention.bindings.push_back({ match.sink->id, 0, match.sink->byte_count }); + attention.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + result.dispatches.push_back(std::move(attention)); + result.covered_nodes.push_back(context.root_index); + return true; +} + static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const FlashAttentionMatch match = match_flash_attention_f32_f16(context.graph, context.plan, context.root_node); @@ -518,30 +664,40 @@ static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & } const CommonActivationPublicationCapabilities capabilities = { - match.query_token_count >= 128 ? DispatchActivationInputF16Row : DispatchActivationInputNone, + match.sink == nullptr && match.query_token_count >= 128 ? DispatchActivationInputF16Row : + DispatchActivationInputNone, DispatchActivationInputF16Row, }; const CommonActivationPublicationPlan publication_plan = common_select_activation_publication_plan(context, *match.output, capabilities); const CommonActivationPublicationDemand * demand = publication_plan.find(CommonActivationPublicationFormat::F16Row); - const bool publish_f16 = demand != nullptr; - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(publish_f16 ? kFlashAttentionF32F16WmmaPublishF16Kernel : - kFlashAttentionF32F16WmmaKernel); + const bool publish_f16 = demand != nullptr; + const bool batch_heads = match.sink != nullptr && !publish_f16 && match.query_token_count <= 4 && + match.key_value_head_count == 1 && match.query_head_count % 4 == 0 && + match.query_head_count >= 4 && + match.query->nb[2] == static_cast(match.qk_head_size) * sizeof(float); + const int64_t logical_tokens = batch_heads ? 4 * match.query_token_count : match.query_token_count; + const int64_t logical_heads = batch_heads ? match.query_head_count / 4 : match.query_head_count; + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.sink != nullptr ? kFlashAttentionF32F16WmmaSinkKernel : + publish_f16 ? kFlashAttentionF32F16WmmaPublishF16Kernel : + kFlashAttentionF32F16WmmaKernel); if (!use_exact_prefill_workload(match.query_token_count, match.key_value_token_count)) { dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; } - dispatch.kernel.integer_parameters.emplace("query_token_count", match.query_token_count); + dispatch.kernel.integer_parameters.emplace("query_token_count", logical_tokens); dispatch.kernel.integer_parameters.emplace("key_value_token_count", match.key_value_token_count); - add_flash_attention_compile_parameters(dispatch.kernel, match.query_head_count, match.key_value_head_count, + add_flash_attention_compile_parameters(dispatch.kernel, logical_heads, match.key_value_head_count, match.qk_head_size, match.value_head_size, match.attention_scale, false, 1, 1); + dispatch.kernel.compile_parameters.emplace("ggml.flash_attention.head_batch_size", batch_heads ? "4" : "1"); dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); dispatch.bindings.push_back(prepare_flash_attention_value(context, match, dispatch_match, dispatch.kernel)); dispatch.bindings.push_back({ match.mask_binding_value, 0, match.mask_binding_bytes }); - dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); + dispatch.bindings.push_back(match.sink ? DispatchBinding{ match.sink->id, 0, match.sink->byte_count } : + DispatchBinding{ match.query->id, 0, match.query->byte_count }); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_f16) { @@ -679,6 +835,22 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa } // namespace void register_flash_attention_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.flash_attention_sparse", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::SingleOp, + 90, + DispatchSource::Common, + match_sparse_attention, + }); + registry.add({ + "common.lightning_indexer_f32_f16", + GGML_OP_LIGHTNING_INDEXER, + DispatchMatchKind::SingleOp, + 50, + DispatchSource::Common, + match_lightning_indexer, + }); registry.add({ "common.flash_attention_f32_f16_gate", GGML_OP_FLASH_ATTN_EXT, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp index 022d1f05279e..bfd078c07b89 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp @@ -19,6 +19,7 @@ namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kGetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32"); +static constexpr KernelCatalogRef kGetRows32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_32"); static constexpr KernelCatalogRef kGetRowsF32NextKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32_next"); static constexpr KernelCatalogRef kGetRowsScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_scale_f32"); @@ -329,6 +330,61 @@ static Dispatch make_get_rows_next_dispatch(const GetRowsMatch & match, return dispatch; } +static bool match_get_rows_32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { + return false; + } + const Value * source = graph_value(context.graph, node->inputs[0]); + const Value * ids = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (source == nullptr || ids == nullptr || output == nullptr || + (source->type != GGML_TYPE_F32 && source->type != GGML_TYPE_I32) || output->type != source->type || + ids->type != GGML_TYPE_I32 || !output->contiguous || source->nb[0] != sizeof(int32_t) || + ids->nb[0] != sizeof(int32_t) || source->ne[2] != ids->ne[1] || source->ne[3] != ids->ne[2] || + ids->ne[3] != 1 || output->ne[0] != source->ne[0] || output->ne[1] != ids->ne[0] || + output->ne[2] != ids->ne[1] || output->ne[3] != ids->ne[2] || output->element_count < 1 || + output->element_count > 134217728 || source->storage == output->storage || ids->storage == output->storage) { + return false; + } + for (const Value * value : { source, ids }) { + if (value->byte_count == 0 || value->byte_count / sizeof(int32_t) > 1073741824 || + value->byte_count % sizeof(int32_t) != 0) { + return false; + } + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (value->ne[dim] < 1 || value->ne[dim] > 1073741824 || value->nb[dim] % sizeof(int32_t) != 0 || + value->nb[dim] / sizeof(int32_t) > 1073741824) { + return false; + } + } + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRows32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + auto & config = dispatch.kernel.compile_parameters; + config.emplace("ggml.get_rows_32.width", std::to_string(source->ne[0])); + config.emplace("ggml.get_rows_32.rows", std::to_string(source->ne[1])); + config.emplace("ggml.get_rows_32.selected", std::to_string(ids->ne[0])); + config.emplace("ggml.get_rows_32.batch1", std::to_string(ids->ne[1])); + config.emplace("ggml.get_rows_32.source_span", std::to_string(source->byte_count / sizeof(int32_t))); + config.emplace("ggml.get_rows_32.ids_span", std::to_string(ids->byte_count / sizeof(int32_t))); + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + config.emplace("ggml.get_rows_32.source_stride" + std::to_string(dim), + std::to_string(source->nb[dim] / sizeof(int32_t))); + } + for (int dim = 1; dim < 3; ++dim) { + config.emplace("ggml.get_rows_32.ids_stride" + std::to_string(dim), + std::to_string(ids->nb[dim] / sizeof(int32_t))); + } + dispatch.bindings.push_back({ source->id, 0, source->byte_count }); + dispatch.bindings.push_back({ ids->id, 0, ids->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + static bool match_get_rows_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const GetRowsMatch match = match_get_rows_f32(context.graph, context.root_node); if (!match.matched()) { @@ -472,6 +528,14 @@ static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & } // namespace void register_get_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.get_rows.32", + GGML_OP_GET_ROWS, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_get_rows_32_dispatch, + }); registry.add({ "common.get_rows_scale.f32", GGML_OP_GET_ROWS, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp index 66f13832a181..f2cb47124238 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp @@ -1,10 +1,12 @@ #include "dispatch-glu.h" +#include "dispatch-common.h" #include "dispatch-layout-utils.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" +#include #include #include #include @@ -108,12 +110,14 @@ static void add_packed_glu_parameters(Dispatch & dispatch, dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(output.ne[0])); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(output.ne[1])); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", std::to_string(output.ne[2])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride0", "1"); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", std::to_string(input.nb[1] / sizeof(float))); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", std::to_string(input.nb[2] / sizeof(float))); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", std::to_string(input.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride0", "1"); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", std::to_string(input.nb[1] / sizeof(float))); dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", @@ -184,9 +188,44 @@ static bool match_packed_glu_f32_dispatch(const DispatchMatchContext & context, return true; } +static bool match_swiglu_clamp_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_GLU || node->inputs.size() != 2) { + return false; + } + const GluParams * params = op_params_as(node->params); + if (params == nullptr || params->op != GGML_GLU_OP_SWIGLU_CLAMP || params->swapped || + !std::isfinite(params->limit) || params->limit < 0.0f) { + return false; + } + const Value * gate = graph_value(context.graph, node->inputs[0]); + const Value * up = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (gate == nullptr || up == nullptr || output == nullptr || gate->type != GGML_TYPE_F32 || + up->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !gate->contiguous || !up->contiguous || + !output->contiguous || gate->ne != output->ne || up->ne != output->ne || output->element_count < 1 || + output->element_count > 134217728 || gate->storage == output->storage || up->storage == output->storage) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_swiglu_clamp_f32")); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.swiglu_clamp.limit", float_config_value(params->limit)); + dispatch.bindings = { + { gate->id, 0, gate->byte_count }, + { up->id, 0, up->byte_count }, + { output->id, 0, output->byte_count } + }; + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + } // namespace void register_glu_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ "common.swiglu_clamp_f32", GGML_OP_GLU, DispatchMatchKind::SingleOp, 10, DispatchSource::Common, + match_swiglu_clamp_f32_dispatch }); registry.add({ "common.packed_glu_f32", GGML_OP_GLU, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.cpp new file mode 100644 index 000000000000..f39846a8a4a7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.cpp @@ -0,0 +1,105 @@ +#include "dispatch-hyper-connection.h" + +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include + +namespace ggml::hrx { +namespace { +static constexpr KernelCatalogRef kComb = GGML_HRX_KERNEL_REF("loom_libs", "ggml_hc_comb_f32"); +static constexpr KernelCatalogRef kPre = GGML_HRX_KERNEL_REF("loom_libs", "ggml_hc_pre_f32"); +static constexpr KernelCatalogRef kPost = GGML_HRX_KERNEL_REF("loom_libs", "ggml_hc_post_f32"); + +static bool match_hyper_connection(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr) { + return false; + } + const bool comb = node->op == GGML_OP_DSV4_HC_COMB; + const bool pre = node->op == GGML_OP_DSV4_HC_PRE; + const bool post = node->op == GGML_OP_DSV4_HC_POST; + const bool identity = post && node->inputs.size() == 3; + if ((!comb && !pre && !post) || node->inputs.size() != (comb ? 3 : pre ? 2 : identity ? 3 : 4)) { + return false; + } + const Value * output = context.graph.values().find(node->output); + if (!output || output->type != GGML_TYPE_F32 || !output->contiguous || output->ne[3] != 1) { + return false; + } + const Value * inputs[4] = {}; + for (size_t i = 0; i < node->inputs.size(); ++i) { + inputs[i] = context.graph.values().find(node->inputs[i]); + if (!inputs[i] || inputs[i]->type != GGML_TYPE_F32 || !inputs[i]->contiguous || inputs[i]->ne[3] != 1 || + inputs[i]->storage == output->storage) { + return false; + } + } + const int64_t tokens = pre ? inputs[0]->ne[2] : inputs[0]->ne[1]; + const int64_t embd = comb ? 1 : inputs[0]->ne[0]; + const int64_t streams = comb ? 4 : pre ? inputs[0]->ne[1] : inputs[1]->ne[1]; + if (tokens < 1 || tokens > 32768 || embd < 1 || embd > 32768 || streams < 1 || streams > 128 || + tokens * embd > 134217728 / streams) { + return false; + } + const auto * pre_params = pre ? op_params_as(node->params) : nullptr; + if (comb) { + if (inputs[0]->ne[0] != 24 || inputs[0]->ne[2] != 1 || inputs[1]->ne[0] < 3 || inputs[1]->ne[1] != 1 || + inputs[1]->ne[2] != 1 || inputs[2]->element_count != 24 || output->ne[0] != 4 || output->ne[1] != 4 || + output->ne[2] != tokens) { + return false; + } + } else if (pre) { + if (pre_params == nullptr || !std::isfinite(pre_params->scale) || output->ne[0] != embd || + output->ne[1] != tokens || output->ne[2] != 1 || + (pre_params->gated ? inputs[1]->ne != inputs[0]->ne : + inputs[1]->ne != std::array{ streams, tokens, 1, 1 })) { + return false; + } + } else if (inputs[0]->ne[2] != 1 || inputs[1]->ne[0] != embd || inputs[1]->ne[2] != tokens || + inputs[2]->ne != std::array{ streams, tokens, 1, 1 } || + (!identity && inputs[3]->ne != std::array{ streams, streams, tokens, 1 }) || + output->ne != inputs[1]->ne) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(comb ? kComb : pre ? kPre : kPost); + dispatch.kernel.compile_parameters.emplace("ggml.hc.tokens", std::to_string(tokens)); + if (comb) { + const auto * params = op_params_as(node->params); + if (!params || !std::isfinite(params->eps) || params->eps <= 0 || params->n_iter < 1 || params->n_iter > 128) { + return false; + } + dispatch.kernel.compile_parameters.emplace("ggml.hc.eps", float_config_value(params->eps)); + dispatch.kernel.compile_parameters.emplace("ggml.hc.n_iter", std::to_string(params->n_iter)); + } else { + dispatch.kernel.compile_parameters.emplace("ggml.hc.embd", std::to_string(embd)); + dispatch.kernel.compile_parameters.emplace("ggml.hc.streams", std::to_string(streams)); + if (pre) { + dispatch.kernel.compile_parameters.emplace("ggml.hc.scale", float_config_value(pre_params->scale)); + dispatch.kernel.compile_parameters.emplace("ggml.hc.gated", pre_params->gated ? "1" : "0"); + } else { + dispatch.kernel.compile_parameters.emplace("ggml.hc.identity", identity ? "1" : "0"); + } + } + for (size_t i = 0; i < node->inputs.size(); ++i) { + dispatch.bindings.push_back({ inputs[i]->id, 0, inputs[i]->byte_count }); + } + if (identity) { + // The identity specialization does not read the mixing matrix. + dispatch.bindings.push_back({ inputs[1]->id, 0, inputs[1]->byte_count }); + } + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} +} // namespace + +void register_hyper_connection_dispatches(DispatchRegistryBuilder & registry) { + for (ggml_op op : { GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, GGML_OP_DSV4_HC_POST }) { + registry.add({ "common.hyper_connection", op, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_hyper_connection }); + } +} +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.h new file mode 100644 index 000000000000..84d83f4f31ff --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-hyper-connection.h @@ -0,0 +1,7 @@ +#pragma once + +#include "dispatch-common.h" + +namespace ggml::hrx { +void register_hyper_connection_dispatches(DispatchRegistryBuilder & registry); +} diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h index e0c6d594732a..4b640e08e384 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h @@ -8,11 +8,13 @@ namespace ggml::hrx { -inline bool strided_f32_storage_span_bytes(const Value & value, size_t & byte_count) { - if (value.type != GGML_TYPE_F32 || value.nb[0] != sizeof(float)) { +inline bool strided_f32_storage_span_bytes(const Value & value, + size_t & byte_count, + bool allow_strided_elements = false) { + if (value.type != GGML_TYPE_F32 || (!allow_strided_elements && value.nb[0] != sizeof(float))) { return false; } - for (int i = 1; i < GGML_MAX_DIMS; ++i) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { if (value.nb[i] % sizeof(float) != 0) { return false; } diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h index 93127a1e6c58..437bcec13898 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h @@ -166,6 +166,8 @@ inline bool common_mul_mat_uses_k16_major_f16(CommonMulMatWeightFormat format, inline ggml_type common_mul_mat_format_type(CommonMulMatWeightFormat format) { switch (format) { + case CommonMulMatWeightFormat::MXFP4: + return GGML_TYPE_MXFP4; case CommonMulMatWeightFormat::Q1_0: return GGML_TYPE_Q1_0; case CommonMulMatWeightFormat::Q3K: @@ -638,7 +640,7 @@ inline bool common_prepare_q8_1_x4_input(const DispatchMatchContext & context, constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); constexpr const char * name = "common.mul_mat.q8_1_x4"; - const ValueId activation = context.next_plan_value; + const ValueId activation(context.next_plan_value.value + static_cast(match.transients.size())); match.transients.push_back({ activation, name, bytes, 256 }); Dispatch quantize; quantize.kernel = make_kernel_specialization(kernel); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h index 741fa22ed615..4d036674bedc 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h @@ -224,7 +224,10 @@ inline CommonMulMatIdMatch common_match_mul_mat_id_any_format(const Graph & !common_mul_mat_id_supported_token_count(token_count) || !common_mul_mat_id_supported_route_count(route_count) || !common_mul_mat_id_supported_expert_count(expert_count) || input_route_count <= 0 || - input_route_count > route_count || route_count % input_route_count != 0) { + input_route_count > route_count || route_count % input_route_count != 0 || + route_ids->nb[0] != sizeof(int32_t) || route_ids->nb[1] % sizeof(int32_t) != 0 || + route_ids->nb[1] < static_cast(route_count) * sizeof(int32_t) || + route_ids->nb[1] / sizeof(int32_t) > 512) { return {}; } @@ -286,8 +289,9 @@ inline bool common_mul_mat_id_ensure_moe_routing_bundle(const DispatchMatchConte const size_t expert_table_bytes = common_mul_mat_id_expert_table_size(match.token_count, match.expert_count); const size_t partition_table_bytes = common_mul_mat_id_partition_table_size(match.token_count, match.route_count, match.expert_count); - const int64_t route_stride = match.route_count; - const size_t route_ids_length = static_cast(match.token_count * route_stride) * sizeof(int32_t); + const int64_t route_stride = match.route_ids->nb[1] / sizeof(int32_t); + const size_t route_ids_length = + static_cast((match.token_count - 1) * route_stride + match.route_count) * sizeof(int32_t); dispatch_match.transients.push_back( { expert_table_value, "common.moe_routing.expert_table", expert_table_bytes, 256 }); 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 5c9822c5aa9e..69042d1b739f 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 @@ -1,5 +1,6 @@ #include "dispatch-mul-mat-id.h" +#include "dispatch-mul-mat-common.h" #include "dispatch-mul-mat-id-common.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -18,6 +19,14 @@ 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 kMulMatIdVectorMxfp4Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_mxfp4_f32"); +static constexpr KernelCatalogRef kMulMatIdVectorMxfp4Q8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_mxfp4_q8"); +static constexpr KernelCatalogRef kMulMatIdPairedMxfp4Q8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_paired_mxfp4_q8"); +static constexpr KernelCatalogRef kMulMatIdMxfp4PrefillPrefetchKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_mxfp4_prefill_prefetch"); static constexpr KernelCatalogRef kMulMatIdTiledPostOpsF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"); static constexpr KernelCatalogRef kMulMatIdSkinnyPostOpsF32F32Kernel = @@ -179,6 +188,66 @@ static MulMatIdPostOpsMatch match_mul_mat_id_postops(const DispatchMatchContext } // namespace +static bool match_mul_mat_id_vector_mxfp4(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatIdMatch match = + common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); + if (!match.matched() || match.weight_format != CommonMulMatWeightFormat::MXFP4 || match.token_count > 5 || + match.input_size % 256 != 0) { + return false; + } + Dispatch dispatch; + const bool integer_dot = match.input_size % 512 == 0; + if (!integer_dot && match.token_count != 1) { + return false; + } + DispatchBinding activation{ match.input->id, 0, match.input->byte_count }; + if (integer_dot && !common_prepare_q8_1_x4_input(context, *match.input, match.input_size, + match.input_route_count * match.token_count, dispatch_match, + activation, CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + return false; + } + const bool pair_tokens = integer_dot && match.token_count > 1 && match.route_count <= 8 && + match.input_size >= 2048 && match.output_size >= 1024; + dispatch.kernel = make_kernel_specialization(pair_tokens ? kMulMatIdPairedMxfp4Q8Kernel : + integer_dot ? kMulMatIdVectorMxfp4Q8Kernel : + kMulMatIdVectorMxfp4Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + const std::pair parameters[] = { + { "input_size", match.input_size }, + { "output_size", match.output_size }, + { "route_count", match.route_count }, + { "input_route_count", match.input_route_count }, + { "route_stride", static_cast(match.route_ids->nb[1] / sizeof(int32_t)) }, + { "expert_count", match.expert_count } + }; + for (const auto & parameter : parameters) { + dispatch.kernel.compile_parameters.emplace(std::string("ggml.mxfp4_vector.") + parameter.first, + std::to_string(parameter.second)); + } + dispatch.bindings.push_back(activation); + dispatch.bindings.push_back({ match.route_ids->id, 0, match.route_ids->byte_count }); + if (pair_tokens) { + const ValueId pairs(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + const size_t bytes = static_cast(match.token_count * match.route_count) * sizeof(int32_t); + dispatch_match.transients.push_back({ pairs, "common.mul_mat_id.pairs", bytes, 256 }); + Dispatch prepare; + prepare.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_mxfp4_prepare_pairs")); + prepare.kernel.integer_parameters = dispatch.kernel.integer_parameters; + prepare.kernel.compile_parameters = dispatch.kernel.compile_parameters; + prepare.bindings = { + { match.route_ids->id, 0, match.route_ids->byte_count }, + { pairs, 0, bytes } + }; + dispatch_match.dispatches.push_back(std::move(prepare)); + dispatch.bindings.push_back({ pairs, 0, bytes }); + } + 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_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + static bool match_mul_mat_id_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const CommonMulMatIdMatch match = common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); @@ -192,7 +261,16 @@ static bool match_mul_mat_id_dispatch(const DispatchMatchContext & context, Disp } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(common_mul_mat_id_f32_kernel(match)); + const bool mxfp4_prefetch = match.weight_format == CommonMulMatWeightFormat::MXFP4 && match.token_count >= 32 && + match.output_size % 64 == 0 && match.input_size * match.input_route_count <= 32768; + DispatchBinding activation{ match.input->id, 0, match.input->byte_count }; + if (mxfp4_prefetch && !common_prepare_q8_1_x4_input( + context, *match.input, match.input_size * match.input_route_count, match.token_count, + dispatch_match, activation, CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + return false; + } + dispatch.kernel = make_kernel_specialization(mxfp4_prefetch ? kMulMatIdMxfp4PrefillPrefetchKernel : + 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)); @@ -209,7 +287,7 @@ static bool match_mul_mat_id_dispatch(const DispatchMatchContext & context, Disp dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_id.weight_format", common_mul_mat_id_to_config_value(common_mul_mat_format_config_value(match.weight_format))); - dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back(activation); dispatch.bindings.push_back({ routing_bundle.expert_table, 0, routing_bundle.expert_table_byte_count }); dispatch.bindings.push_back({ routing_bundle.partition_table, 0, routing_bundle.partition_table_byte_count }); dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); @@ -312,6 +390,14 @@ static bool match_mul_mat_id_postops_dispatch(const DispatchMatchContext & conte } void register_mul_mat_id_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.mul_mat_id.vector_mxfp4", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::SingleOp, + 125, + DispatchSource::Common, + match_mul_mat_id_vector_mxfp4, + }); registry.add({ "common.mul_mat_id_postops.f32_f32_wmma", GGML_OP_MUL_MAT_ID, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h index cc3848c42e10..6402038c265b 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h @@ -29,6 +29,7 @@ enum class CommonMulMatWeightFormat { IQ4_XS, Q8_0, Q8_1, + MXFP4, F16, BF16, F32, @@ -90,6 +91,9 @@ inline bool common_mul_mat_format_for_type(ggml_type type, CommonMulMatWeightFor case GGML_TYPE_Q8_1: format = CommonMulMatWeightFormat::Q8_1; return true; + case GGML_TYPE_MXFP4: + format = CommonMulMatWeightFormat::MXFP4; + return true; case GGML_TYPE_F16: format = CommonMulMatWeightFormat::F16; return true; @@ -150,6 +154,8 @@ inline int64_t common_mul_mat_format_config_value(CommonMulMatWeightFormat forma return 80; case CommonMulMatWeightFormat::Q8_1: return 81; + case CommonMulMatWeightFormat::MXFP4: + return 39; case CommonMulMatWeightFormat::F16: return 16; case CommonMulMatWeightFormat::BF16: @@ -176,6 +182,7 @@ inline bool common_mul_mat_supported_dense_input_size(CommonMulMatWeightFormat f case CommonMulMatWeightFormat::IQ4_NL: case CommonMulMatWeightFormat::Q8_0: case CommonMulMatWeightFormat::Q8_1: + case CommonMulMatWeightFormat::MXFP4: return input_size >= 256 && input_size <= 32768 && input_size % 32 == 0; default: return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; 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 4e785ac0fb79..d0206537fbc1 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 @@ -1,6 +1,7 @@ #include "dispatch-mul-mat.h" #include "dispatch-activation-publication.h" +#include "dispatch-copy.h" #include "dispatch-mul-mat-common.h" #include "dispatch_registration/immutable-program-resource.h" #include "ggml.h" @@ -17,6 +18,9 @@ namespace ggml::hrx { namespace { +static constexpr KernelCatalogRef kMulMatVectorQ8Q8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q8_q8_1"); + static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatTiledF32F32Kernel = @@ -1295,12 +1299,15 @@ static bool build_mul_mat_dispatch( const bool use_q1_token64 = use_aligned_dynamic && match.token_count == 64 && match.weight_format == CommonMulMatWeightFormat::Q1_0 && std::getenv("GGML_HRX_DISABLE_Q1_TOKEN64_MATMUL") == nullptr; + const bool use_q8_token64 = + use_aligned_dynamic && match.token_count % 64 == 0 && match.weight_format == CommonMulMatWeightFormat::Q8_0; Dispatch dispatch; - dispatch.kernel = make_kernel_specialization( - use_iq_codebook ? tiled_iq_codebook_kernel(match.weight_format) : - use_q1_token64 ? kMulMatTiledToken64AlignedF32F32Kernel : - use_aligned_dynamic ? kMulMatTiledAlignedF32F32Kernel : - publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : match.kernel); + dispatch.kernel = + make_kernel_specialization(use_iq_codebook ? tiled_iq_codebook_kernel(match.weight_format) : + (use_q1_token64 || use_q8_token64) ? kMulMatTiledToken64AlignedF32F32Kernel : + use_aligned_dynamic ? kMulMatTiledAlignedF32F32Kernel : + publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : + match.kernel); if (use_aligned_dynamic) { dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; } @@ -1576,6 +1583,31 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, } } + const auto * precision = op_params_as(context.root_node->params); + if (match.weight_format == CommonMulMatWeightFormat::Q8_0 && match.input_size % 256 == 0 && + match.output_size >= 512 && !match.has_fused_unary && + (precision == nullptr || precision->prec == GGML_PREC_DEFAULT) && + vector_publish_format_config_value(publish) == 0) { + DispatchBinding quantized; + if (!common_prepare_q8_1_x4_input(context, *match.input, match.input_size, 1, dispatch_match, quantized, + CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + return false; + } + dispatch.kernel = make_kernel_specialization(kMulMatVectorQ8Q8Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", 1); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", std::to_string(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", std::to_string(match.output_size)); + dispatch.bindings = { + quantized, + { match.weight->id, 0, match.weight->byte_count }, + { match.output->id, 0, match.output->byte_count }, + { 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 append_vector_q8_publish(*match.output, 1, match.output_size, publish, dispatch_match); + } + DispatchBinding activation; CommonMulMatWeightFormat weight_format = match.weight_format; const char * weight_layout = nullptr; @@ -2022,6 +2054,228 @@ static bool match_skinny_mul_mat_dispatch(const DispatchMatchContext & context, return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); } +static bool match_batched_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return false; + } + const Value * weight = common_graph_value(context.graph, node->inputs[0]); + const Value * input = common_graph_value(context.graph, node->inputs[1]); + const Value * output = common_graph_value(context.graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr || !weight->contiguous || !output->contiguous || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + (weight->type != GGML_TYPE_Q8_0 && weight->type != GGML_TYPE_F16 && weight->type != GGML_TYPE_F32) || + weight->ne[2] <= 1 || weight->ne[2] > 512 || weight->ne[3] != 1 || input->ne[3] != 1 || output->ne[3] != 1 || + weight->ne[2] != input->ne[2] || output->ne[2] != input->ne[2] || weight->storage == output->storage || + input->storage == output->storage) { + return false; + } + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + CommonMulMatWeightFormat format; + if (!common_mul_mat_format_for_type(weight->type, format) || + !common_is_supported_dense_input_size(format, input_size) || + !common_is_supported_dense_output_size(output_size) || token_count < 1 || token_count > 2048 || + input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count) { + return false; + } + + ValueId input_value = input->id; + if (!input->contiguous) { + Value packed = *input; + packed.id = context.next_plan_value; + packed.storage = ValueStorageId(); + packed.contiguous = true; + packed.byte_count = packed.element_count * sizeof(float); + size_t stride = sizeof(float); + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + packed.nb[dim] = stride; + stride *= packed.ne[dim]; + } + Dispatch copy; + if (!make_copy_f32_dispatch(*input, packed, copy)) { + return false; + } + dispatch_match.transients.push_back({ packed.id, "common.mul_mat.batched_input", packed.byte_count, 256 }); + dispatch_match.dispatches.push_back(std::move(copy)); + input_value = packed.id; + } + + const size_t input_bytes = input_size * token_count * sizeof(float); + const size_t output_bytes = output_size * token_count * sizeof(float); + const size_t weight_bytes = weight->nb[2]; + const bool vector = token_count == 1; + const auto * precision = op_params_as(node->params); + const bool q8_vector = vector && format == CommonMulMatWeightFormat::Q8_0 && input_size % 256 == 0 && + output_size >= 512 && (precision == nullptr || precision->prec == GGML_PREC_DEFAULT); + DispatchBinding quantized; + if (q8_vector) { + Value packed = *input; + packed.id = input_value; + packed.contiguous = true; + if (!common_prepare_q8_1_x4_input(context, packed, input_size, weight->ne[2], dispatch_match, quantized, + CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatVectorQ8Q8Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", 1); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", std::to_string(input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", std::to_string(output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.batch_count", std::to_string(weight->ne[2])); + dispatch.bindings = { + quantized, + { weight->id, 0, weight->byte_count }, + { output->id, 0, output->byte_count }, + { output->id, 0, output->byte_count } + }; + dispatch_match.dispatches.push_back(std::move(dispatch)); + dispatch_match.covered_nodes.push_back(context.root_index); + return true; + } + const bool q8_token64 = !vector && token_count % 64 == 0 && format == CommonMulMatWeightFormat::Q8_0; + const std::string prefix = vector ? "ggml.matmul.vector." : "ggml.mul_mat."; + for (int64_t batch = 0; batch < weight->ne[2]; ++batch) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(vector ? kMulMatVectorF32F32Kernel : + q8_token64 ? kMulMatTiledToken64AlignedF32F32Kernel : + kMulMatTiledF32F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", std::to_string(token_count)); + dispatch.kernel.compile_parameters.emplace(prefix + "input_size", std::to_string(input_size)); + dispatch.kernel.compile_parameters.emplace(prefix + "output_size", std::to_string(output_size)); + dispatch.kernel.compile_parameters.emplace(prefix + "weight_format", + std::to_string(common_mul_mat_format_config_value(format))); + dispatch.kernel.compile_parameters.emplace(prefix + "output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace(prefix + "output_unary_op", "23"); + dispatch.kernel.compile_parameters.emplace(prefix + "activation_format", "0"); + dispatch.bindings.push_back({ input_value, batch * input_bytes, input_bytes }); + dispatch.bindings.push_back({ weight->id, batch * weight_bytes, weight_bytes }); + dispatch.bindings.push_back({ output->id, batch * output_bytes, output_bytes }); + if (vector) { + dispatch.bindings.push_back({ output->id, batch * output_bytes, output_bytes }); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + } + dispatch_match.covered_nodes.push_back(context.root_index); + return true; +} + +static bool match_small_m_vector(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return false; + } + const Value * weight = context.graph.values().find(node->inputs[0]); + const Value * input = context.graph.values().find(node->inputs[1]); + const Value * output = context.graph.values().find(node->output); + if (!weight || !input || !output || weight->type != GGML_TYPE_F32 || input->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !weight->contiguous || !input->contiguous || !output->contiguous || + weight->ne[0] < 4096 || weight->ne[0] > 32768 || weight->ne[0] % 2048 != 0 || weight->ne[1] < 1 || + weight->ne[1] > 32 || weight->ne[2] != 1 || weight->ne[3] != 1 || + input->ne != std::array{ weight->ne[0], 1, 1, 1 } || + output->ne != std::array{ weight->ne[1], 1, 1, 1 } || output->storage == input->storage || + output->storage == weight->storage) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_small_m_f32")); + dispatch.kernel.compile_parameters.emplace("ggml.small_m.k", std::to_string(weight->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.small_m.m", std::to_string(weight->ne[1])); + dispatch.bindings = { + { weight->id, 0, weight->byte_count }, + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +static bool match_hadamard_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return false; + } + const auto * params = op_params_as(node->params); + if (params == nullptr || params->hint != GGML_HINT_SRC0_IS_HADAMARD) { + return false; + } + const Value * weight = context.graph.values().find(node->inputs[0]); + const Value * input = context.graph.values().find(node->inputs[1]); + const Value * output = context.graph.values().find(node->output); + if (weight == nullptr || input == nullptr || output == nullptr || weight->type != GGML_TYPE_F32 || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !input->contiguous || !output->contiguous || + input->ne != output->ne || weight->ne != std::array{ 128, 128, 1, 1 } || input->ne[0] != 128 || + input->element_count < 128 || input->element_count > 134217728 || input->storage == output->storage) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_hadamard_128_f32")); + dispatch.kernel.integer_parameters.emplace("row_count", input->element_count / 128); + dispatch.bindings = { + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +static bool match_small_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return false; + } + const Value * weight = context.graph.values().find(node->inputs[0]); + const Value * input = context.graph.values().find(node->inputs[1]); + const Value * output = context.graph.values().find(node->output); + if (weight == nullptr || input == nullptr || output == nullptr || + (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) || input->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !output->contiguous || weight->ne[0] < 1 || weight->ne[0] > 256 || + input->ne[0] != weight->ne[0] || output->ne[0] != weight->ne[1] || output->ne[1] != input->ne[1] || + output->element_count < 1 || output->element_count > 134217728 || output->storage == input->storage || + output->storage == weight->storage) { + return false; + } + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (weight->ne[dim] < 1 || input->ne[dim] < 1 || weight->nb[dim] % ggml_type_size(weight->type) != 0 || + input->nb[dim] % sizeof(float) != 0 || weight->nb[dim] == 0 || input->nb[dim] == 0 || + (dim > 1 && (input->ne[dim] % weight->ne[dim] != 0 || output->ne[dim] != input->ne[dim]))) { + return false; + } + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_small_k_f32")); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + auto & config = dispatch.kernel.compile_parameters; + auto set = [&](const std::string & key, int64_t value) { + config.emplace("ggml.small_matmul." + key, std::to_string(value)); + }; + set("k", weight->ne[0]); + set("m", output->ne[0]); + set("n", output->ne[1]); + set("ne2", output->ne[2]); + set("repeat2", input->ne[2] / weight->ne[2]); + set("repeat3", input->ne[3] / weight->ne[3]); + set("weight_format", weight->type == GGML_TYPE_F16 ? 16 : 32); + set("weight_span", weight->byte_count / ggml_type_size(weight->type)); + set("input_span", input->byte_count / sizeof(float)); + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + set("ws" + std::to_string(dim), weight->nb[dim] / ggml_type_size(weight->type)); + set("xs" + std::to_string(dim), input->nb[dim] / sizeof(float)); + } + dispatch.bindings = { + { weight->id, 0, weight->byte_count }, + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + 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); @@ -2111,6 +2365,10 @@ static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & conte } void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ "common.mul_mat.hadamard128", GGML_OP_MUL_MAT, DispatchMatchKind::SingleOp, 350, + DispatchSource::Common, match_hadamard_dispatch }); + registry.add({ "common.mul_mat.small_k", GGML_OP_MUL_MAT, DispatchMatchKind::SingleOp, -10, DispatchSource::Common, + match_small_mul_mat_dispatch }); registry.add_activation_consumer({ "common.mul_mat.input.f16_row.factored_tiled", GGML_OP_MUL_MAT, @@ -2194,6 +2452,14 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, common_accepts_f16_row_prefill_mul_mat, }); + registry.add({ + "common.mul_mat.vector_small_m_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 320, + DispatchSource::Common, + match_small_m_vector, + }); registry.add({ "common.mul_mat.vector_q6_selected_outputs", GGML_OP_MUL_MAT, @@ -2330,6 +2596,14 @@ void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Common, match_tiled_mul_mat_dispatch, }); + registry.add({ + "common.mul_mat.batched", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 80, + DispatchSource::Common, + match_batched_mul_mat_dispatch, + }); registry.add({ "common.mul_mat.skinny_f32_f32", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.cpp new file mode 100644 index 000000000000..94edc1fa4058 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.cpp @@ -0,0 +1,235 @@ +#include "dispatch-reduction.h" + +#include "dispatch-common.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { +constexpr KernelCatalogRef kSumRowsF32 = GGML_HRX_KERNEL_REF("loom_libs", "ggml_sum_rows_f32"); +constexpr KernelCatalogRef kSoftMaxF32 = GGML_HRX_KERNEL_REF("loom_libs", "ggml_softmax_f32"); + +bool match_argmax(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_ARGMAX || node->inputs.size() != 1) { + return false; + } + const Value * input = context.graph.values().find(node->inputs[0]); + const Value * output = context.graph.values().find(node->output); + if (input == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_I32 || + !input->contiguous || !output->contiguous || input->ne[0] < 1 || input->ne[0] > 1048576 || input->ne[1] < 1 || + input->ne[1] > 2048 || input->ne[2] != 1 || input->ne[3] != 1 || output->element_count != input->ne[1] || + input->storage == output->storage) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_argmax_f32")); + dispatch.kernel.integer_parameters.emplace("row_count", input->ne[1]); + dispatch.kernel.compile_parameters.emplace("ggml.argmax.width", std::to_string(input->ne[0])); + dispatch.bindings = { + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +bool match_top_k_radix(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_TOP_K || node->inputs.size() != 1) { + return false; + } + const Value * input = context.graph.values().find(node->inputs[0]); + const Value * output = context.graph.values().find(node->output); + if (input == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_I32 || + !input->contiguous || !output->contiguous || input->ne[0] <= 512 || input->ne[0] > 131072 || + output->ne[0] < 1 || output->ne[0] > 512 || output->ne[0] > input->ne[0] || input->storage == output->storage) { + return false; + } + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + if (input->ne[dim] != output->ne[dim]) { + return false; + } + } + const int64_t rows = input->element_count / input->ne[0]; + if (rows < 1 || rows > 2048) { + return false; + } + const int64_t count = output->ne[0]; + if (rows <= 16 && input->ne[0] >= 4096 && count <= 32) { + int64_t width = input->ne[0]; + ValueId ids = input->id; + size_t ids_bytes = input->byte_count; + int stage = 0; + do { + const int64_t next_width = width / 256 * count + std::min(count, width % 256); + const bool last = next_width == count; + const size_t bytes = static_cast(rows * next_width) * sizeof(int32_t); + const ValueId target = last ? output->id : ValueId(context.next_plan_value.value + stage); + if (!last) { + match.transients.push_back({ target, "common.top_k.partitions", bytes, 256 }); + } + Dispatch dispatch; + dispatch.kernel = + make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_top_k_radix_partitions_f32")); + dispatch.kernel.integer_parameters.emplace("row_count", rows); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.width", std::to_string(input->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.count", std::to_string(count)); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.stage_width", std::to_string(width)); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.first_pass", stage == 0 ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.last_pass", last ? "1" : "0"); + dispatch.bindings = { + { input->id, 0, input->byte_count }, + { ids, 0, ids_bytes }, + { target, 0, bytes } + }; + match.dispatches.push_back(std::move(dispatch)); + ids = target; + ids_bytes = bytes; + width = next_width; + ++stage; + } while (width > count); + match.covered_nodes.push_back(context.root_index); + return true; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_top_k_radix_f32")); + dispatch.kernel.integer_parameters.emplace("row_count", rows); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.width", std::to_string(input->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.top_k.count", std::to_string(output->ne[0])); + dispatch.bindings = { + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +bool match_argsort(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || (node->op != GGML_OP_ARGSORT && node->op != GGML_OP_TOP_K) || node->inputs.size() != 1) { + return false; + } + const bool top_k = node->op == GGML_OP_TOP_K; + const auto * params = op_params_as(node->params); + const Value * input = context.graph.values().find(node->inputs[0]); + const Value * output = context.graph.values().find(node->output); + if ((!top_k && params == nullptr) || input == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_I32 || !input->contiguous || !output->contiguous || + (!top_k && input->ne != output->ne) || output->ne[0] < 1 || output->ne[0] > input->ne[0] || input->ne[0] < 1 || + input->ne[0] > 2048 || input->element_count / input->ne[0] > 1048576 || input->storage == output->storage) { + return false; + } + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + if (input->ne[dim] != output->ne[dim]) { + return false; + } + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_argsort_f32")); + dispatch.kernel.integer_parameters.emplace("row_count", input->element_count / input->ne[0]); + dispatch.kernel.compile_parameters.emplace("ggml.argsort.width", std::to_string(input->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.argsort.descending", + top_k || params->order == GGML_SORT_ORDER_DESC ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace("ggml.argsort.output_width", std::to_string(output->ne[0])); + dispatch.bindings = { + { input->id, 0, input->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +bool match_softmax(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_SOFT_MAX || node->inputs.size() != 1) { + return false; + } + const auto * params = op_params_as(node->params); + const Value * input = context.graph.values().find(node->inputs[0]); + const Value * output = context.graph.values().find(node->output); + if (params == nullptr || input == nullptr || output == nullptr || !std::isfinite(params->scale) || + params->max_bias != 0.0f || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input->contiguous || !output->contiguous || input->ne != output->ne || input->ne[0] < 1 || + input->ne[0] > 32768 || input->element_count < 1 || input->element_count > 134217728 || + (input->storage == output->storage && input->storage_offset != output->storage_offset)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kSoftMaxF32); + dispatch.kernel.integer_parameters.emplace("row_count", input->element_count / input->ne[0]); + dispatch.kernel.compile_parameters.emplace("ggml.softmax.width", std::to_string(input->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.softmax.scale", float_config_value(params->scale)); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + +bool match_sum_rows(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || (node->op != GGML_OP_SUM_ROWS && node->op != GGML_OP_SUM) || node->inputs.size() != 1) { + return false; + } + const bool full_sum = node->op == GGML_OP_SUM; + const Value * input = context.graph.values().find(node->inputs[0]); + const Value * output = context.graph.values().find(node->output); + if (input == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !output->contiguous || input->ne[0] < 1 || (!full_sum && input->ne[0] > 32768) || output->ne[0] != 1 || + (full_sum && (!input->contiguous || input->element_count > 16777216 || output->element_count != 1)) || + output->element_count < 1 || output->element_count > 1048576 || input->storage == output->storage || + input->byte_count == 0 || input->byte_count % sizeof(float) != 0 || + input->byte_count / sizeof(float) > 1073741824) { + return false; + } + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (input->ne[dim] < 1 || input->nb[dim] % sizeof(float) != 0 || input->nb[dim] / sizeof(float) > 1073741824 || + (!full_sum && dim > 0 && input->ne[dim] != output->ne[dim])) { + return false; + } + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kSumRowsF32); + dispatch.kernel.integer_parameters.emplace("row_count", output->element_count); + auto & config = dispatch.kernel.compile_parameters; + config.emplace("ggml.sum_rows.width", std::to_string(full_sum ? input->element_count : input->ne[0])); + config.emplace("ggml.sum_rows.ne1", std::to_string(full_sum ? 1 : input->ne[1])); + config.emplace("ggml.sum_rows.ne2", std::to_string(full_sum ? 1 : input->ne[2])); + config.emplace("ggml.sum_rows.source_span", std::to_string(input->byte_count / sizeof(float))); + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + config.emplace("ggml.sum_rows.stride" + std::to_string(dim), std::to_string(input->nb[dim] / sizeof(float))); + } + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} +} // namespace + +void register_reduction_dispatches(DispatchRegistryBuilder & registry) { + registry.add( + { "common.argmax.f32", GGML_OP_ARGMAX, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, match_argmax }); + registry.add({ "common.top_k.radix.f32", GGML_OP_TOP_K, DispatchMatchKind::SingleOp, 20, DispatchSource::Common, + match_top_k_radix }); + registry.add( + { "common.top_k.f32", GGML_OP_TOP_K, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, match_argsort }); + registry.add( + { "common.sum.f32", GGML_OP_SUM, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, match_sum_rows }); + registry.add({ "common.argsort.f32", GGML_OP_ARGSORT, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_argsort }); + registry.add({ "common.softmax.f32", GGML_OP_SOFT_MAX, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_softmax }); + registry.add({ "common.sum_rows.f32", GGML_OP_SUM_ROWS, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_sum_rows }); +} +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.h new file mode 100644 index 000000000000..d2ca3c9c9030 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-reduction.h @@ -0,0 +1,6 @@ +#pragma once +#include "dispatch_registration/dispatch-registry.h" + +namespace ggml::hrx { +void register_reduction_dispatches(DispatchRegistryBuilder & registry); +} 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 a1976d98e4e5..c424ec17f2ce 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 @@ -315,7 +315,8 @@ struct RopeConcatMatch { static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Status * status = nullptr) { RopeMatch match; - if (node == nullptr || node->op != GGML_OP_ROPE || node->inputs.size() < 2 || node->inputs.size() > 3) { + if (node == nullptr || (node->op != GGML_OP_ROPE && node->op != GGML_OP_ROPE_BACK) || node->inputs.size() < 2 || + node->inputs.size() > 3) { log_rope_reject(status, "root is not a supported ROPE arity", nullptr, nullptr, nullptr); return match; } @@ -360,6 +361,14 @@ static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Sta log_rope_reject(status, "failed to build supported ROPE theta table", input, positions, output); return {}; } + if (node->op == GGML_OP_ROPE_BACK) { + for (size_t offset = 0; offset < theta_data.size(); offset += sizeof(float)) { + float theta; + std::memcpy(&theta, theta_data.data() + offset, sizeof(theta)); + theta = -theta; + std::memcpy(theta_data.data() + offset, &theta, sizeof(theta)); + } + } const Value * freq_factors = nullptr; size_t freq_factors_bytes = 0; @@ -401,7 +410,10 @@ static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Sta return match; } -static SetRowsMatch match_set_rows_2d(const Graph & graph, const GraphNode * node, Status * status = nullptr) { +static SetRowsMatch match_set_rows_2d(const Graph & graph, + const GraphNode * node, + Status * status = nullptr, + bool allow_i32_indices = false) { SetRowsMatch match; if (node == nullptr || node->op != GGML_OP_SET_ROWS || node->inputs.size() != 3) { log_set_rows_reject(status, "root is not a supported SET_ROWS arity", nullptr, nullptr, nullptr, nullptr); @@ -413,8 +425,10 @@ static SetRowsMatch match_set_rows_2d(const Graph & graph, const GraphNode * nod const Value * cache = graph_value(graph, node->inputs[2]); const Value * output = graph_value(graph, node->output); if (rows == nullptr || indices == nullptr || cache == nullptr || output == nullptr || !indices->contiguous || - !cache->contiguous || indices->type != GGML_TYPE_I64 || output->type != cache->type || - !same_shape(*cache, *output) || !graph.values().same_storage(cache->id, output->id)) { + !cache->contiguous || + (indices->type != GGML_TYPE_I64 && !(allow_i32_indices && indices->type == GGML_TYPE_I32)) || + output->type != cache->type || !same_shape(*cache, *output) || + !graph.values().same_storage(cache->id, output->id)) { log_set_rows_reject(status, "input/output shape or storage is unsupported", rows, indices, cache, output); return {}; } @@ -489,6 +503,8 @@ static void add_set_rows_compile_parameters(Dispatch & dispatch, const SetRowsMa dispatch.kernel.compile_parameters.emplace("ggml.set_rows.input_format", to_config_value(match.row_format)); dispatch.kernel.compile_parameters.emplace("ggml.set_rows.output_format", to_config_value(match.output_format)); dispatch.kernel.compile_parameters.emplace("ggml.set_rows.input_stride", to_config_value(match.input_stride)); + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.index_format", + match.indices->type == GGML_TYPE_I32 ? "32" : "64"); } static ValueId add_constant_binding(const DispatchMatchContext & context, @@ -858,7 +874,7 @@ static bool match_rope_f32_dispatch(const DispatchMatchContext & context, Dispat } static bool match_set_rows_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const SetRowsMatch match = match_set_rows_2d(context.graph, context.root_node, &dispatch_match.status); + const SetRowsMatch match = match_set_rows_2d(context.graph, context.root_node, &dispatch_match.status, true); if (!match.matched()) { return false; } @@ -870,7 +886,77 @@ static bool match_set_rows_dispatch(const DispatchMatchContext & context, Dispat } // namespace +static bool match_set_rows_scalar(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_SET_ROWS || node->inputs.size() != 3) { + return false; + } + const Value * rows = graph_value(context.graph, node->inputs[0]); + const Value * ids = graph_value(context.graph, node->inputs[1]); + const Value * cache = graph_value(context.graph, node->inputs[2]); + const Value * output = graph_value(context.graph, node->output); + if (rows == nullptr || ids == nullptr || cache == nullptr || output == nullptr || + (rows->type != GGML_TYPE_F32 && rows->type != GGML_TYPE_F16) || rows->type != cache->type || + cache->type != output->type || (ids->type != GGML_TYPE_I32 && ids->type != GGML_TYPE_I64) || rows->ne[0] != 1 || + cache->ne[0] != 1 || ids->ne[0] != rows->ne[1] || ids->ne[3] != 1 || rows->ne[2] != cache->ne[2] || + rows->ne[3] != cache->ne[3] || cache->ne != output->ne || rows->element_count < 1 || + rows->element_count > 134217728 || cache->storage != output->storage || + cache->storage_offset != output->storage_offset || rows->storage == output->storage || + ids->storage == output->storage) { + return false; + } + for (const Value * value : { rows, ids, cache }) { + const size_t bytes = ggml_type_size(value->type); + if (value->byte_count / bytes > 1073741824) { + return false; + } + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (value->ne[dim] < 1 || value->nb[dim] == 0 || value->nb[dim] % bytes != 0 || + value->nb[dim] / bytes > 1073741824) { + return false; + } + } + } + if (rows->ne[2] % ids->ne[1] != 0 || rows->ne[3] % ids->ne[2] != 0) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_set_rows_scalar")); + dispatch.kernel.integer_parameters.emplace("element_count", rows->element_count); + auto & config = dispatch.kernel.compile_parameters; + auto set = [&](const std::string & key, int64_t value) { + config.emplace("ggml.set_rows_scalar." + key, std::to_string(value)); + }; + set("nr", rows->ne[1]); + set("ne2", rows->ne[2]); + set("in1", ids->ne[1]); + set("in2", ids->ne[2]); + set("cache_rows", cache->ne[1]); + set("index_bits", ggml_type_size(ids->type) * 8); + set("element_bits", ggml_type_size(rows->type) * 8); + set("source_span", rows->byte_count / ggml_type_size(rows->type)); + set("index_span", ids->byte_count / ggml_type_size(ids->type)); + set("cache_span", cache->byte_count / ggml_type_size(cache->type)); + for (int dim = 0; dim < 3; ++dim) { + set("is" + std::to_string(dim), ids->nb[dim] / ggml_type_size(ids->type)); + set("rs" + std::to_string(dim + 1), rows->nb[dim + 1] / ggml_type_size(rows->type)); + set("cs" + std::to_string(dim + 1), cache->nb[dim + 1] / ggml_type_size(cache->type)); + } + dispatch.bindings = { + { rows->id, 0, rows->byte_count }, + { ids->id, 0, ids->byte_count }, + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + void register_rope_set_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ "common.rope_back.f32", GGML_OP_ROPE_BACK, DispatchMatchKind::SingleOp, 100, DispatchSource::Common, + match_rope_f32_dispatch }); + registry.add({ "common.set_rows.scalar", GGML_OP_SET_ROWS, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, + match_set_rows_scalar }); registry.add({ "common.rope_scale.f32", GGML_OP_ROPE, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp index 03cbc2b8b232..f63a960791de 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp @@ -1,13 +1,12 @@ #include "dispatch-unary.h" +#include "dispatch-common.h" #include "ggml.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include #include -#include #include -#include #include #include @@ -15,7 +14,9 @@ namespace ggml::hrx { namespace { static constexpr KernelCatalogRef kUnaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_unary_f32"); +static constexpr KernelCatalogRef kSoftPlusF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_softplus_f32"); static constexpr KernelCatalogRef kScaleBiasF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_bias_f32"); +static constexpr KernelCatalogRef kClampF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_clamp_f32"); static bool same_shape(const Value & lhs, const Value & rhs) { for (int i = 0; i < GGML_MAX_DIMS; ++i) { @@ -30,12 +31,6 @@ static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); } -static std::string float_config_value(float value) { - std::ostringstream stream; - stream << std::setprecision(std::numeric_limits::max_digits10) << value; - return stream.str(); -} - static bool scale_storage_is_safe(const Value & input, const Value & output) { if (input.storage_root != output.storage_root) { return true; @@ -61,7 +56,7 @@ static bool match_scale_bias_f32_dispatch(const DispatchMatchContext & context, const Value * input = graph_value(context.graph, node->inputs[0]); const Value * output = graph_value(context.graph, node->output); if (params == nullptr || input == nullptr || output == nullptr || !std::isfinite(params->scale) || - !std::isfinite(params->bias) || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + std::isnan(params->bias) || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !input->contiguous || !output->contiguous || !same_shape(*input, *output) || !scale_storage_is_safe(*input, *output) || output->element_count <= 0 || static_cast(output->element_count) > std::numeric_limits::max()) { @@ -88,7 +83,7 @@ static bool match_unary_f32_dispatch(const DispatchMatchContext & context, Dispa } const UnaryParams * params = op_params_as(node->params); - if (params == nullptr || !unary_kind_supported(params->op)) { + if (params == nullptr || (!unary_kind_supported(params->op) && params->op != UnaryKind::SoftPlus)) { return false; } @@ -105,11 +100,14 @@ static bool match_unary_f32_dispatch(const DispatchMatchContext & context, Dispa } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kUnaryF32Kernel); + dispatch.kernel = + make_kernel_specialization(params->op == UnaryKind::SoftPlus ? kSoftPlusF32Kernel : kUnaryF32Kernel); dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); - dispatch.kernel.compile_parameters.emplace("ggml.unary_f32.op", - std::to_string(unary_kind_config_value(params->op))); + if (params->op != UnaryKind::SoftPlus) { + dispatch.kernel.compile_parameters.emplace("ggml.unary_f32.op", + std::to_string(unary_kind_config_value(params->op))); + } dispatch.bindings.push_back({ input->id, 0, input->byte_count }); dispatch.bindings.push_back({ output->id, 0, output->byte_count }); @@ -118,6 +116,57 @@ static bool match_unary_f32_dispatch(const DispatchMatchContext & context, Dispa return true; } +static bool match_clamp_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->inputs.size() != 1) { + return false; + } + const auto * params = op_params_as(node->params); + const Value * input = graph_value(context.graph, node->inputs[0]); + const Value * output = graph_value(context.graph, node->output); + if (params == nullptr || input == nullptr || output == nullptr || std::isnan(params->min) || + std::isnan(params->max) || params->min > params->max || input->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !input->contiguous || !output->contiguous || !same_shape(*input, *output) || + !scale_storage_is_safe(*input, *output) || output->element_count < 1 || output->element_count > 134217728) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kClampF32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.clamp.min", float_config_value(params->min)); + dispatch.kernel.compile_parameters.emplace("ggml.clamp.max", float_config_value(params->max)); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_fill(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_FILL) { + return false; + } + const auto * params = op_params_as(node->params); + const Value * output = context.graph.values().find(node->output); + if (params == nullptr || output == nullptr || std::isnan(params->value) || !output->contiguous || + (output->type != GGML_TYPE_F32 && output->type != GGML_TYPE_F16) || output->element_count < 1 || + output->element_count > 134217728) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(GGML_HRX_KERNEL_REF("loom_libs", "ggml_fill")); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.fill.value", float_config_value(params->value)); + dispatch.kernel.compile_parameters.emplace("ggml.fill.format", output->type == GGML_TYPE_F16 ? "16" : "32"); + dispatch.bindings = { + { output->id, 0, output->byte_count } + }; + match.dispatches.push_back(std::move(dispatch)); + match.covered_nodes.push_back(context.root_index); + return true; +} + static void register_unary_dispatch_for(DispatchRegistryBuilder & registry, ggml_op root_op) { registry.add({ "common.unary_f32", @@ -132,6 +181,15 @@ static void register_unary_dispatch_for(DispatchRegistryBuilder & registry, ggml } // namespace void register_unary_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ "common.fill", GGML_OP_FILL, DispatchMatchKind::SingleOp, 0, DispatchSource::Common, match_fill }); + registry.add({ + "common.clamp_f32", + GGML_OP_CLAMP, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_clamp_f32_dispatch, + }); registry.add({ "common.scale_bias_f32", GGML_OP_SCALE, diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp index 7555e200fa60..e08981afc67f 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -3,6 +3,7 @@ #include "backend-buffer-binding.h" #include "common/dispatch-common.h" +#include "graph/graph-traversal.h" #include "llm/dispatch-attention-qkv.h" #include "llm/dispatch-gated-delta-net.h" #include "llm/dispatch-ssm-conv.h" @@ -25,6 +26,33 @@ static bool valid_root_op(ggml_op op) { return op >= 0 && static_cast(op) < kOpCount; } +static bool match_inputs_are_available(const DispatchMatchContext & context, const DispatchMatch & match) { + for (size_t node_index : match.covered_nodes) { + if (node_index >= context.graph.nodes().size()) { + return false; + } + for (ValueId input : context.graph.nodes()[node_index].inputs) { + const GraphNode * producer = context.graph.index().producer(input); + while (producer != nullptr) { + size_t index; + if (!context.graph.index().node_index(producer, index)) { + return false; + } + if ((index < context.covered_nodes.size() && context.covered_nodes[index]) || + std::find(match.covered_nodes.begin(), match.covered_nodes.end(), index) != + match.covered_nodes.end()) { + break; + } + if (!is_layout_alias_node(context.graph, *producer)) { + return false; + } + producer = context.graph.index().producer(producer->inputs[0]); + } + } + } + return true; +} + static void sort_registrations(std::vector & registrations) { std::stable_sort( registrations.begin(), registrations.end(), @@ -98,7 +126,8 @@ bool DispatchRegistry::match(const DispatchMatchContext & context, for (const DispatchRegistration & registration : registrations) { DispatchMatch candidate; const bool matched = registration.matcher != nullptr && registration.matcher(context, candidate); - const bool eligible = matched && dispatch_has_eligible_weight_layouts(context, candidate); + const bool eligible = matched && match_inputs_are_available(context, candidate) && + dispatch_has_eligible_weight_layouts(context, candidate); if (eligible) { if (diagnostics != nullptr) { diagnostics->attempts.push_back({ diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 6be3802203f0..e775c46a8d53 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -635,8 +635,7 @@ static void device_props(ggml_backend_dev_t device, ggml_backend_dev_props * pro device_memory(device, &props->memory_free, &props->memory_total); props->type = GGML_BACKEND_DEVICE_TYPE_GPU; props->device_id = nullptr; - props->caps = { true, true, true, false }; - props->caps.mmap_support = true; + props->caps = { true, true, true, false, true }; } static ggml_backend_t device_init(ggml_backend_dev_t device, const char * parameters) { @@ -1068,12 +1067,6 @@ static bool supported_binary_f32_tensor(const ggml_tensor * op) { return false; } - if ((!ggml_are_same_shape(op->src[0], op) || !ggml_are_same_shape(op->src[1], op)) && - (!ggml_is_contiguous(op->src[0]) || !ggml_is_contiguous(op->src[1]) || - !tensor_has_packed_f32_layout(op->src[0]) || !tensor_has_packed_f32_layout(op->src[1]))) { - return false; - } - if (!binary_kind_allows_broadcast(binary_kind, op->src[0], op->src[1], op)) { return false; } @@ -1140,7 +1133,6 @@ static bool supported_unary_f32_tensor(const ggml_tensor * op) { } static bool device_supports_op(ggml_backend_dev_t device, const ggml_tensor * op) { - GGML_UNUSED(device); if (op == nullptr) { return false; } @@ -1164,6 +1156,14 @@ static bool device_supports_op(ggml_backend_dev_t device, const ggml_tensor * op return source_f32 && output_f32 && source_contiguous && output_contiguous && same_shape && separate_output && nonempty && bounded; } + if (op->op == GGML_OP_DSV4_HC_COMB || op->op == GGML_OP_DSV4_HC_PRE || op->op == GGML_OP_DSV4_HC_POST || + op->op == GGML_OP_CLAMP || op->op == GGML_OP_FILL || op->op == GGML_OP_SUM || op->op == GGML_OP_REPEAT || + op->op == GGML_OP_ROPE_BACK || op->op == GGML_OP_TOP_K || op->op == GGML_OP_ARGSORT || + op->op == GGML_OP_ARGMAX || op->op == GGML_OP_CONCAT || op->op == GGML_OP_LIGHTNING_INDEXER || + op->op == GGML_OP_MUL_MAT || (op->op == GGML_OP_ADD && op->type == GGML_TYPE_F16) || + (op->op == GGML_OP_GLU && ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_CLAMP)) { + return ggml::hrx::can_execute_standalone_op_as_graph(op, device_context(device)->architecture); + } const bool supported_binary = supported_binary_f32_tensor(op); if (supported_binary) { return true; diff --git a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp index 2ead551ffe26..47e775ef1565 100644 --- a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp +++ b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp @@ -4,6 +4,7 @@ #include "ggml.h" #include +#include #include #include #include @@ -119,6 +120,7 @@ json op_params_json(const OpParams & params) { { "max_bias", value.max_bias }, { "logit_softcap", value.logit_softcap }, { "prec", static_cast(value.prec) }, + { "n_kv_max", value.n_kv_max }, }; } else if constexpr (std::is_same_v) { return { @@ -126,6 +128,12 @@ json op_params_json(const OpParams & params) { { "scale", value.scale }, { "max_bias", value.max_bias } }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "mul_mat" }, + { "prec", static_cast(value.prec) }, + { "hint", static_cast(value.hint) } + }; } else if constexpr (std::is_same_v) { return { { "kind", "argsort" }, @@ -137,11 +145,29 @@ json op_params_json(const OpParams & params) { { "min", value.min }, { "max", value.max } }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "fill" }, + { "value", std::isfinite(value.value) ? json(value.value) : json(value.value < 0 ? "-inf" : "inf") } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "hc_comb" }, + { "eps", value.eps }, + { "n_iter", value.n_iter } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "hc_pre" }, + { "scale", value.scale }, + { "gated", value.gated } + }; } else if constexpr (std::is_same_v) { return { { "kind", "glu" }, { "op", static_cast(value.op) }, - { "swapped", value.swapped } + { "swapped", value.swapped }, + { "limit", value.limit } }; } else if constexpr (std::is_same_v) { return { @@ -189,11 +215,16 @@ OpParams parse_op_params(const json & item) { item.value("max_bias", 0.0f), item.value("logit_softcap", 0.0f), static_cast(item.value("prec", static_cast(GGML_PREC_DEFAULT))), + item.value("n_kv_max", 0), }; } if (kind == "soft_max") { return SoftMaxParams{ item.value("scale", 0.0f), item.value("max_bias", 0.0f) }; } + if (kind == "mul_mat") { + return MulMatParams{ static_cast(item.value("prec", int(GGML_PREC_DEFAULT))), + static_cast(item.value("hint", int(GGML_HINT_NONE))) }; + } if (kind == "argsort") { return ArgsortParams{ static_cast( item.value("order", static_cast(GGML_SORT_ORDER_ASC))) }; @@ -204,9 +235,19 @@ OpParams parse_op_params(const json & item) { std::numeric_limits::infinity(); return ClampParams{ min, max }; } + if (kind == "fill") { + const auto & value = item.at("value"); + return FillParams{ value.is_string() ? std::stof(value.get()) : value.get() }; + } + if (kind == "hc_comb") { + return HyperConnectionCombParams{ item.value("eps", 0.0f), item.value("n_iter", 0) }; + } + if (kind == "hc_pre") { + return HyperConnectionPreParams{ item.value("scale", 1.0f), item.value("gated", false) }; + } if (kind == "glu") { return GluParams{ static_cast(item.value("op", static_cast(GGML_GLU_OP_REGLU))), - item.value("swapped", false) }; + item.value("swapped", false), item.value("limit", 0.0f) }; } if (kind == "scale") { return ScaleParams{ item.value("scale", 0.0f), item.value("bias", 0.0f) }; diff --git a/ggml/src/ggml-hrx/graph/op-params.cpp b/ggml/src/ggml-hrx/graph/op-params.cpp index bf3374f973bf..0d6c299606d1 100644 --- a/ggml/src/ggml-hrx/graph/op-params.cpp +++ b/ggml/src/ggml-hrx/graph/op-params.cpp @@ -25,7 +25,8 @@ static bool flash_attn_ext_params_equivalent(const OpParams & lhs, const OpParam const FlashAttnExtParams * rhs_params = op_params_as(rhs); return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->scale, rhs_params->scale) && nearly_equal(lhs_params->max_bias, rhs_params->max_bias) && - nearly_equal(lhs_params->logit_softcap, rhs_params->logit_softcap) && lhs_params->prec == rhs_params->prec; + nearly_equal(lhs_params->logit_softcap, rhs_params->logit_softcap) && lhs_params->prec == rhs_params->prec && + lhs_params->n_kv_max == rhs_params->n_kv_max; } static bool soft_max_params_equivalent(const OpParams & lhs, const OpParams & rhs) { @@ -48,11 +49,29 @@ static bool clamp_params_equivalent(const OpParams & lhs, const OpParams & rhs) nearly_equal(lhs_params->max, rhs_params->max); } +static bool fill_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const auto * l = op_params_as(lhs); + const auto * r = op_params_as(rhs); + return l != nullptr && r != nullptr && l->value == r->value; +} + static bool glu_params_equivalent(const OpParams & lhs, const OpParams & rhs) { const GluParams * lhs_params = op_params_as(lhs); const GluParams * rhs_params = op_params_as(rhs); return lhs_params != nullptr && rhs_params != nullptr && lhs_params->op == rhs_params->op && - lhs_params->swapped == rhs_params->swapped; + lhs_params->swapped == rhs_params->swapped && lhs_params->limit == rhs_params->limit; +} + +static bool hyper_connection_comb_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const auto * l = op_params_as(lhs); + const auto * r = op_params_as(rhs); + return l != nullptr && r != nullptr && l->eps == r->eps && l->n_iter == r->n_iter; +} + +static bool hyper_connection_pre_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const auto * l = op_params_as(lhs); + const auto * r = op_params_as(rhs); + return l != nullptr && r != nullptr && l->scale == r->scale && l->gated == r->gated; } static bool scale_params_equivalent(const OpParams & lhs, const OpParams & rhs) { @@ -257,6 +276,9 @@ OpParams import_op_params(const ggml_tensor & tensor) { } switch (tensor.op) { + case GGML_OP_MUL_MAT: + return MulMatParams{ static_cast(ggml_get_op_params_i32(&tensor, 0)), + static_cast(ggml_get_op_params_i32(&tensor, 1)) }; case GGML_OP_RMS_NORM: case GGML_OP_L2_NORM: return RmsNormParams{ ggml_get_op_params_f32(&tensor, 0) }; @@ -267,26 +289,34 @@ OpParams import_op_params(const ggml_tensor & tensor) { }; case GGML_OP_FLASH_ATTN_EXT: return FlashAttnExtParams{ - ggml_get_op_params_f32(&tensor, 0), - ggml_get_op_params_f32(&tensor, 1), - ggml_get_op_params_f32(&tensor, 2), - ggml_flash_attn_ext_get_prec(&tensor), + ggml_get_op_params_f32(&tensor, 0), ggml_get_op_params_f32(&tensor, 1), + ggml_get_op_params_f32(&tensor, 2), ggml_flash_attn_ext_get_prec(&tensor), + ggml_get_op_params_i32(&tensor, 4), }; case GGML_OP_ARGSORT: return ArgsortParams{ static_cast(ggml_get_op_params_i32(&tensor, 0)) }; + case GGML_OP_DSV4_HC_COMB: + return HyperConnectionCombParams{ ggml_get_op_params_f32(&tensor, 0), ggml_get_op_params_i32(&tensor, 1) }; + case GGML_OP_DSV4_HC_PRE: + return HyperConnectionPreParams{ ggml_get_op_params_f32(&tensor, 0), + ggml_get_op_params_i32(&tensor, 1) != 0 }; + case GGML_OP_FILL: + return FillParams{ ggml_get_op_params_f32(&tensor, 0) }; case GGML_OP_CLAMP: return ClampParams{ ggml_get_op_params_f32(&tensor, 0), ggml_get_op_params_f32(&tensor, 1), }; case GGML_OP_GLU: - return GluParams{ ggml_get_glu_op(&tensor), ggml_get_op_params_i32(&tensor, 1) != 0 }; + return GluParams{ ggml_get_glu_op(&tensor), ggml_get_op_params_i32(&tensor, 1) != 0, + ggml_get_op_params_f32(&tensor, 3) }; case GGML_OP_SCALE: return ScaleParams{ ggml_get_op_params_f32(&tensor, 0), ggml_get_op_params_f32(&tensor, 1), }; case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: return RopeParams{ ggml_get_op_params_i32(&tensor, 1), ggml_get_op_params_i32(&tensor, 2), @@ -311,6 +341,11 @@ OpParams import_op_params(const ggml_tensor & tensor) { } bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs) { + if (op == GGML_OP_MUL_MAT) { + const auto * l = op_params_as(lhs); + const auto * r = op_params_as(rhs); + return l != nullptr && r != nullptr && l->prec == r->prec && l->hint == r->hint; + } switch (op) { case GGML_OP_RMS_NORM: case GGML_OP_L2_NORM: @@ -321,6 +356,12 @@ bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs return flash_attn_ext_params_equivalent(lhs, rhs); case GGML_OP_ARGSORT: return argsort_params_equivalent(lhs, rhs); + case GGML_OP_DSV4_HC_COMB: + return hyper_connection_comb_params_equivalent(lhs, rhs); + case GGML_OP_DSV4_HC_PRE: + return hyper_connection_pre_params_equivalent(lhs, rhs); + case GGML_OP_FILL: + return fill_params_equivalent(lhs, rhs); case GGML_OP_CLAMP: return clamp_params_equivalent(lhs, rhs); case GGML_OP_GLU: @@ -340,6 +381,7 @@ bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs case GGML_OP_COS: return unary_params_equivalent(lhs, rhs); case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: return rope_params_equivalent(lhs, rhs); default: return lhs.index() == rhs.index(); diff --git a/ggml/src/ggml-hrx/graph/op-params.h b/ggml/src/ggml-hrx/graph/op-params.h index 73f5a9122322..d06807f175e0 100644 --- a/ggml/src/ggml-hrx/graph/op-params.h +++ b/ggml/src/ggml-hrx/graph/op-params.h @@ -14,11 +14,17 @@ struct RmsNormParams { float eps = 0.0f; }; +struct MulMatParams { + ggml_prec prec = GGML_PREC_DEFAULT; + ggml_op_hint hint = GGML_HINT_NONE; +}; + struct FlashAttnExtParams { float scale = 0.0f; float max_bias = 0.0f; float logit_softcap = 0.0f; ggml_prec prec = GGML_PREC_DEFAULT; + int32_t n_kv_max = 0; }; struct SoftMaxParams { @@ -35,9 +41,24 @@ struct ClampParams { float max = 0.0f; }; +struct FillParams { + float value = 0.0f; +}; + +struct HyperConnectionCombParams { + float eps = 0.0f; + int32_t n_iter = 0; +}; + +struct HyperConnectionPreParams { + float scale = 1.0f; + bool gated = false; +}; + struct GluParams { ggml_glu_op op = GGML_GLU_OP_REGLU; bool swapped = false; + float limit = 0.0f; }; struct ScaleParams { @@ -116,10 +137,14 @@ struct RopeParams { using OpParams = std::variant< std::monostate, RmsNormParams, + MulMatParams, FlashAttnExtParams, SoftMaxParams, ArgsortParams, ClampParams, + FillParams, + HyperConnectionCombParams, + HyperConnectionPreParams, GluParams, ScaleParams, BinaryParams, diff --git a/ggml/src/ggml-hrx/hrx-minimum-commit.txt b/ggml/src/ggml-hrx/hrx-minimum-commit.txt index 3c1fd22364e4..f2c1dd3cb66a 100644 --- a/ggml/src/ggml-hrx/hrx-minimum-commit.txt +++ b/ggml/src/ggml-hrx/hrx-minimum-commit.txt @@ -1 +1 @@ -771711d49befd397f10c3c6deb66f1c5c89e5042 +d03575e84f602d3e39b87895a30fae214f308558 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 625075d1fbc9..cc19da4358b9 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 @@ -652,6 +652,87 @@ }, { "path": "motifs/matmul/tiled_pair_postprocess_f32.loom" + }, + { + "path": "ops/elementwise/clamp_f32.loom" + }, + { + "path": "ops/elementwise/swiglu_clamp_f32.loom" + }, + { + "path": "ops/elementwise/hc_pre_f32.loom" + }, + { + "path": "ops/elementwise/hc_post_f32.loom" + }, + { + "path": "ops/elementwise/hc_comb_f32.loom" + }, + { + "path": "ops/rows/get_rows_32.loom" + }, + { + "path": "ops/rows/sum_rows_f32.loom" + }, + { + "path": "ops/rows/softmax_f32.loom" + }, + { + "path": "ops/matmul/small_k_f32.loom" + }, + { + "path": "ops/rows/argsort_f32.loom" + }, + { + "path": "ops/copy/copy_bits.loom" + }, + { + "path": "ops/rows/set_rows_scalar.loom" + }, + { + "path": "ops/elementwise/add_f16.loom" + }, + { + "path": "ops/copy/concat_packed.loom" + }, + { + "path": "ops/elementwise/fill.loom" + }, + { + "path": "ops/matmul_id/vector_mxfp4_f32.loom" + }, + { + "path": "ops/matmul/hadamard_128_f32.loom" + }, + { + "path": "ops/matmul/vector_q8_q8_1.loom" + }, + { + "path": "ops/matmul_id/vector_mxfp4_q8.loom" + }, + { + "path": "ops/flash_attention/lightning_indexer_f32_f16.loom" + }, + { + "path": "ops/matmul/vector_small_m_f32.loom" + }, + { + "path": "ops/rows/argmax_f32.loom" + }, + { + "path": "ops/matmul_id/mxfp4_prefill_prefetch.loom" + }, + { + "path": "motifs/matmul_id/mxfp4_prefill_prefetch.loom" + }, + { + "path": "ops/flash_attention/compact_mask_f16.loom" + }, + { + "path": "ops/rows/top_k_radix_f32.loom" + }, + { + "path": "ops/flash_attention/sparse_prefill_f32_f16.loom" } ], "exports": [ @@ -9578,6 +9659,1244 @@ "motifs/bf16_f16.loom", "motifs/f32_f16.loom" ] + }, + { + "name": "ggml_clamp_f32", + "family": "loom_libs", + "symbol": "ggml_clamp_f32", + "source": "ops/elementwise/clamp_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/clamp_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_swiglu_clamp_f32", + "family": "loom_libs", + "symbol": "ggml_swiglu_clamp_f32", + "source": "ops/elementwise/swiglu_clamp_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "gate", + "up", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/swiglu_clamp_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_hc_pre_f32", + "family": "loom_libs", + "symbol": "ggml_hc_pre_f32", + "source": "ops/elementwise/hc_pre_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "weights", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/hc_pre_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_hc_post_f32", + "family": "loom_libs", + "symbol": "ggml_hc_post_f32", + "source": "ops/elementwise/hc_post_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "residual", + "post", + "comb", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/hc_post_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_hc_comb_f32", + "family": "loom_libs", + "symbol": "ggml_hc_comb_f32", + "source": "ops/elementwise/hc_comb_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "mixes", + "scale", + "base", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/hc_comb_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_flash_attention_f32_f16_wmma_sink", + "family": "loom_libs", + "symbol": "ggml_flash_attention_f32_f16_wmma_sink", + "source": "ops/flash_attention/prefill_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "sink", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + { + "name": "ggml_softplus_f32", + "family": "loom_libs", + "symbol": "ggml_softplus_f32", + "source": "ops/elementwise/unary_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/unary_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/unary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/unary_apply.loom" + ] + }, + { + "name": "ggml_get_rows_32", + "family": "loom_libs", + "symbol": "ggml_get_rows_32", + "source": "ops/rows/get_rows_32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "source", + "ids", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_sum_rows_f32", + "family": "loom_libs", + "symbol": "ggml_sum_rows_f32", + "source": "ops/rows/sum_rows_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/sum_rows_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_softmax_f32", + "family": "loom_libs", + "symbol": "ggml_softmax_f32", + "source": "ops/rows/softmax_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/softmax_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_small_k_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_small_k_f32", + "source": "ops/matmul/small_k_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "weight", + "input", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/small_k_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_argsort_f32", + "family": "loom_libs", + "symbol": "ggml_argsort_f32", + "source": "ops/rows/argsort_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/argsort_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_copy_bits", + "family": "loom_libs", + "symbol": "ggml_copy_bits", + "source": "ops/copy/copy_bits.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + }, + { + "name": "source_span", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + }, + { + "name": "source_span", + "type": "index" + } + ], + "bindings": [ + "source", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/copy_bits.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_set_rows_scalar", + "family": "loom_libs", + "symbol": "ggml_set_rows_scalar", + "source": "ops/rows/set_rows_scalar.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "rows", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/set_rows_scalar.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_add_f16", + "family": "loom_libs", + "symbol": "ggml_add_f16", + "source": "ops/elementwise/add_f16.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/add_f16.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_concat_packed", + "family": "loom_libs", + "symbol": "ggml_concat_packed", + "source": "ops/copy/concat_packed.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/concat_packed.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_fill", + "family": "loom_libs", + "symbol": "ggml_fill", + "source": "ops/elementwise/fill.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "output" + ], + "binding_access": [ + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/fill.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_id_vector_mxfp4_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_mxfp4_f32", + "source": "ops/matmul_id/vector_mxfp4_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_mxfp4_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_hadamard_128_f32", + "family": "loom_libs", + "symbol": "ggml_hadamard_128_f32", + "source": "ops/matmul/hadamard_128_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/hadamard_128_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_vector_q8_q8_1", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_q8_q8_1", + "source": "ops/matmul/vector_q8_q8_1.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", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_q8_q8_1.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_id_vector_mxfp4_q8", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_mxfp4_q8", + "source": "ops/matmul_id/vector_mxfp4_q8.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_mxfp4_q8.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mxfp4_prepare_pairs", + "family": "loom_libs", + "symbol": "ggml_mxfp4_prepare_pairs", + "source": "ops/matmul_id/vector_mxfp4_q8.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "ids", + "pairs" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_mxfp4_q8.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_id_paired_mxfp4_q8", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_paired_mxfp4_q8", + "source": "ops/matmul_id/vector_mxfp4_q8.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "ids", + "pairs", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_mxfp4_q8.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_lightning_indexer_f32_f16", + "family": "loom_libs", + "symbol": "ggml_lightning_indexer_f32_f16", + "source": "ops/flash_attention/lightning_indexer_f32_f16.loom", + "workload_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "weight", + "mask", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/lightning_indexer_f32_f16.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_vector_small_m_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_small_m_f32", + "source": "ops/matmul/vector_small_m_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_small_m_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_argmax_f32", + "family": "loom_libs", + "symbol": "ggml_argmax_f32", + "source": "ops/rows/argmax_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/argmax_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_id_mxfp4_prefill_prefetch", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_mxfp4_prefill_prefetch", + "source": "ops/matmul_id/mxfp4_prefill_prefetch.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/mxfp4_prefill_prefetch.loom" + ], + "library_sources": [ + "motifs/matmul_id/mxfp4_prefill_prefetch.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/mxfp4_prefill_prefetch.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_compact_mask_f16", + "family": "loom_libs", + "symbol": "ggml_compact_mask_f16", + "source": "ops/flash_attention/compact_mask_f16.loom", + "workload_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "bindings": [ + "mask", + "indices", + "compact_mask" + ], + "binding_access": [ + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/compact_mask_f16.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_gather_sparse_kv_f16", + "family": "loom_libs", + "symbol": "ggml_gather_sparse_kv_f16", + "source": "ops/flash_attention/compact_mask_f16.loom", + "workload_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_count", + "type": "index" + }, + { + "name": "key_count", + "type": "index" + } + ], + "bindings": [ + "indices", + "key", + "compact_key" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/compact_mask_f16.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_top_k_radix_f32", + "family": "loom_libs", + "symbol": "ggml_top_k_radix_f32", + "source": "ops/rows/top_k_radix_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/top_k_radix_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_top_k_radix_partitions_f32", + "family": "loom_libs", + "symbol": "ggml_top_k_radix_partitions_f32", + "source": "ops/rows/top_k_radix_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + } + ], + "bindings": [ + "input", + "input_ids", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/top_k_radix_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_sparse_prefill_f32_f16", + "family": "loom_libs", + "symbol": "ggml_sparse_prefill_f32_f16", + "source": "ops/flash_attention/sparse_prefill_f32_f16.loom", + "workload_parameters": [ + { + "name": "query_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "mask", + "sink", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/sparse_prefill_f32_f16.loom" + ], + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] } ], "link_modules": [ @@ -9640,6 +10959,5 @@ "source": "ops/rows/get_rows_scale_f32.loom" } ], - "sources": [], - "kernels": [] + "sources": [] } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom index d4964444b253..02d2b0316ca9 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom @@ -2649,6 +2649,58 @@ func.def inline @ggml_f32_f32_vector4(%weight: buffer, %row_byte_base: offset, % func.return %values : vector<4xf32> } +// GGUF stores sixteen low nibbles, then sixteen high nibbles, per scale byte. +func.def inline @ggml_mxfp4_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %block_bytes = index.constant 17 : offset + %one_byte = index.constant 1 : offset + %block = index.div %k, %c32 : index + %block_offset = index.scale %block, %block_bytes : index, offset -> offset + %scale_offset = index.add %row_byte_base, %block_offset : offset + %payload_offset = index.add %scale_offset, %one_byte : offset + %byte_index = index.rem %k, %c16 : index + %within_block = index.rem %k, %c32 : index + %half = index.div %within_block, %c16 : index + %shift_index = index.mul %half, %c4 : index + %shift = index.cast %shift_index : index to i32 + %scale_view = buffer.view %weight[%scale_offset] : buffer -> view<1xi8> + %payload_view = buffer.view %weight[%payload_offset] : buffer -> view<16xi8> + %scale_byte = view.load %scale_view[%c0] : view<1xi8> -> i8 + %scale_word = scalar.extui %scale_byte : i8 to i32 + %bytes = vector.load %payload_view[%byte_index] : view<16xi8> -> vector<4xi8> + %bits = vector.bitcast %bytes : vector<4xi8> to vector<1xi32> + %word = vector.extract %bits[%c0] : vector<1xi32> -> i32 + %selected = scalar.shrui %word, %shift : i32 + %payload = vector.splat %selected : vector<1xi32> + %schema = encoding.define #encoding.f4e2m1 : encoding + %decoded = vector.decode %payload using %schema : vector<1xi32>, encoding -> vector<8xf16> + %v0 = vector.extract %decoded[%c0] : vector<8xf16> -> f16 + %v1 = vector.extract %decoded[%c2] : vector<8xf16> -> f16 + %v2 = vector.extract %decoded[%c4] : vector<8xf16> -> f16 + %v3 = vector.extract %decoded[%c6] : vector<8xf16> -> f16 + %values_half = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf16> + %values = vector.extf %values_half : vector<4xf16> to vector<4xf32> + %zero = scalar.constant 0 : i32 + %nan_exponent = scalar.constant 255 : i32 + %exponent_shift = scalar.constant 23 : i32 + %subnormal_bits = scalar.constant 4194304 : i32 + %nan_bits = scalar.constant 2143289344 : i32 + %normal_bits = scalar.shli %scale_word, %exponent_shift : i32 + %is_zero = scalar.cmpi eq, %scale_word, %zero : i32 + %is_nan = scalar.cmpi eq, %scale_word, %nan_exponent : i32 + %finite_bits = scf.select %is_zero, %subnormal_bits, %normal_bits : i32 + %scale_bits = scf.select %is_nan, %nan_bits, %finite_bits : i32 + %scale = scalar.bitcast %scale_bits : i32 to f32 + %scales = vector.splat %scale : vector<4xf32> + %result = vector.mulf %values, %scales : vector<4xf32> + func.return %result : vector<4xf32> +} + func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) { %q1_0_format = index.constant 10 : index %iq1_s_format = index.constant 19 : index @@ -2673,6 +2725,7 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %q5_1_format = index.constant 51 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index + %mxfp4_format = index.constant 39 : index %zero_bytes = index.constant 0 : offset %q1_0_tile_bytes = index.constant 36 : offset %q3_tile_bytes = index.constant 110 : offset @@ -2717,6 +2770,7 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %is_iq2_xs = index.cmp eq, %weight_format, %c25 : index %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index + %is_mxfp4 = index.cmp eq, %weight_format, %mxfp4_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index %is_f32 = index.cmp eq, %weight_format, %c32 : index @@ -2743,7 +2797,9 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %selected_f16 = scf.select %is_f16, %f16_tile_bytes, %selected_q8_1 : offset %selected_bf16 = scf.select %is_bf16, %bf16_tile_bytes, %selected_f16 : offset %selected = scf.select %is_f32, %f32_tile_bytes, %selected_bf16 : offset - func.return %selected : offset + %mxfp4_bytes = index.constant 136 : offset + %result = scf.select %is_mxfp4, %mxfp4_bytes, %selected : offset + func.return %result : offset } func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) { @@ -2770,6 +2826,7 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi %q5_1_format = index.constant 51 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index + %mxfp4_format = index.constant 39 : index %c32 = index.constant 32 : index %c128 = index.constant 128 : index %c256 = index.constant 256 : index @@ -2820,6 +2877,7 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi %is_iq4_nl = index.cmp eq, %weight_format, %c20 : index %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index + %is_mxfp4 = index.cmp eq, %weight_format, %mxfp4_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index %is_f32 = index.cmp eq, %weight_format, %c32_format : index @@ -2847,7 +2905,10 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi %selected_f16 = scf.select %is_f16, %f16_row_bytes, %selected_q8_1 : offset %selected_bf16 = scf.select %is_bf16, %f16_row_bytes, %selected_f16 : offset %selected = scf.select %is_f32, %f32_row_bytes, %selected_bf16 : offset - func.return %selected : offset + %mxfp4_block_bytes = index.constant 17 : offset + %mxfp4_row_bytes = index.scale %legacy_block_count, %mxfp4_block_bytes : index, offset -> offset + %result = scf.select %is_mxfp4, %mxfp4_row_bytes, %selected : offset + func.return %result : offset } // Centralizes row-format decoding for FP16 WMMA staging. Each format path is @@ -2880,6 +2941,7 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %q5_1_format = index.constant 51 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index + %mxfp4_format = index.constant 39 : index %zero = vector.constant 0.0 : vector<4xf16> // Decode the compact integer format id into one predicate per layout. %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index @@ -2897,6 +2959,7 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %is_iq4_xs = index.cmp eq, %weight_format, %c23 : index %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index + %is_mxfp4 = index.cmp eq, %weight_format, %mxfp4_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index %is_f32 = index.cmp eq, %weight_format, %c32 : index @@ -3037,7 +3100,15 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf16> %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf16> %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf16> - func.return %selected : vector<4xf16> + %mxfp4_values = scf.if %is_mxfp4 -> (vector<4xf16>) { + %decoded = func.call @ggml_mxfp4_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %decoded : vector<4xf32> to vector<4xf16> + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %result = scf.select %is_mxfp4, %mxfp4_values, %selected : vector<4xf16> + func.return %result : vector<4xf16> } // Centralizes row-format decoding for F32 row-gather outputs. @@ -3061,6 +3132,7 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %q5_1_format = index.constant 51 : index %q8_0_format = index.constant 80 : index %q8_1_format = index.constant 81 : index + %mxfp4_format = index.constant 39 : index %zero = vector.constant 0.0 : vector<4xf32> %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index %is_q3 = index.cmp eq, %weight_format, %c3 : index @@ -3077,6 +3149,7 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %is_iq4_xs = index.cmp eq, %weight_format, %c23 : index %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index + %is_mxfp4 = index.cmp eq, %weight_format, %mxfp4_format : index %is_f16 = index.cmp eq, %weight_format, %c16 : index %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index %is_f32 = index.cmp eq, %weight_format, %c32 : index @@ -3208,7 +3281,14 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf32> %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf32> %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf32> - func.return %selected : vector<4xf32> + %mxfp4_values = scf.if %is_mxfp4 -> (vector<4xf32>) { + %decoded = func.call @ggml_mxfp4_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %result = scf.select %is_mxfp4, %mxfp4_values, %selected : vector<4xf32> + func.return %result : vector<4xf32> } // Reuse immutable Q5 metadata across all four 64-element quarters. diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_f32.loom index 8b74883e14d0..48d1aa26e088 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_f32.loom @@ -4,13 +4,13 @@ config.decl @ggml.flash_attention.gate_stride_head : %value: index where [range( config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range(%value, 1, 16777216)] -template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer, %output_offset: offset, %f16_output_offset: offset) template.decl @ggml.flash_attention.decode_split.publish_f32.scalar(%query_head_count: index, %value_head_size: index, %query_head: index, %output_channel: index, %sum: f32, %unnormalized_output: f32, %output: buffer) template.decl @ggml.flash_attention.decode_split.publish_f32.vector2(%query_head_count: index, %value_head_size: index, %query_head_valid0: i1, %query_head_valid1: i1, %lane_output_channel_valid: i1, %query_head0: index, %query_head1: index, %lane_output_channel: index, %sum0: f32, %sum1: f32, %unnormalized_output0: vector<2xf32>, %unnormalized_output1: vector<2xf32>, %output: buffer) -template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_attention_prefill_publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) { +template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_attention_prefill_publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer, %output_offset: offset, %f16_output_offset: offset) { %c0_offset = index.constant 0 : offset %c1 = index.constant 1 : index %c1_f32 = scalar.constant 1.0 : f32 @@ -21,8 +21,8 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at %apply_gate = index.cmp eq, %apply_gate_index, %c1 : index %gate_aligned = buffer.assume.alignment %gate {minimum_alignment = 16} : buffer %gate_view = buffer.view %gate_aligned[%c0_offset] : buffer -> view<268435456xf32> - %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> - %f16_output_view = buffer.view %f16_output[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + %output_view = buffer.view %output[%output_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + %f16_output_view = buffer.view %f16_output[%f16_output_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> // Normalize and publish the lane-owned packets. Query-tail rows never // participate in the output store. scf.if %lane_has_output { diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom index 9008a1dce89a..308d61104a44 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom @@ -1,23 +1,68 @@ -template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) +config.decl @ggml.flash_attention.value_layout : index + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer, %value_offset: offset) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) -template.def<@ggml.flash_attention.pv_block_wmma.prefill_full> device @ggml_flash_attention_pv_block_wmma_prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { +template.def<@ggml.flash_attention.pv_block_wmma.prefill_full> device @ggml_flash_attention_pv_block_wmma_prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer, %value_offset: offset) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { %c0 = index.constant 0 : index %c1 = index.constant 1 : index %c4 = index.constant 4 : index %c16 = index.constant 16 : index %c64 = index.constant 64 : index %c0_offset = index.constant 0 : offset - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index %probability_layout = encoding.layout.strided [72, 1] : encoding %value_layout = encoding.layout.strided [%value_row_stride, %value_column_stride] : encoding - %value_view = buffer.view %value[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> + %value_view = buffer.view %value[%value_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf32> + // Reuse score scratch after probability publication. + // Zero V when every query has zero probability to avoid signed-zero WMMA rounding. + %visibility_offset = index.constant 4096 : offset + %staged_value_offset = index.constant 2048 : offset + %visibility = buffer.view %product_stage[%visibility_offset] : buffer -> view<64x4xf16> + %staged_values = buffer.view %product_stage[%staged_value_offset] : buffer -> view<16x64xf16> + %zero_value = scalar.constant 0.0 : f16 + %workitem = kernel.workitem.id : index + %lane = kernel.subgroup.lane.id : index + %subgroup = kernel.subgroup.id : index + %c256 = index.constant 256 : index + %p0 = view.load %probability_stage_view[%query_row0, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p1 = view.load %probability_stage_view[%query_row1, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p2 = view.load %probability_stage_view[%query_row2, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p3 = view.load %probability_stage_view[%query_row3, %lane] : view<16x64xf16, %probability_layout> -> f16 + %probabilities = vector.from_elements %p0, %p1, %p2, %p3 : vector<4xf16> + %maximum_probability = vector.reduce %probabilities, %zero_value : vector<4xf16>, f16 + view.store %maximum_probability, %visibility[%lane, %subgroup] : f16, view<64x4xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %value_layout_tag = config.get @ggml.flash_attention.value_layout : index + %direct_value_layout = index.cmp eq, %value_layout_tag, %c0 : index + %tile_visibility_offset = index.constant 4608 : offset + %tile_visibility = buffer.view %product_stage[%tile_visibility_offset] : buffer -> view<4xf32> + %zero_flag = scalar.constant 0.0 : f32 + scf.if %direct_value_layout { + %lane_probabilities = vector.load %visibility[%lane, %c0] : view<64x4xf16> -> vector<4xf16> + %lane_probability = vector.reduce %lane_probabilities, %zero_value : vector<4xf16>, f16 + %lane_key_visible = scalar.cmpf ogt, %lane_probability, %zero_value : f16 + %lane_key_tile = index.div %lane, %c16 : index + // Transposed direct fragments fail instrumented validation; retain staged loads. + %one_flag = scalar.constant 1.0 : f32 + %is_first_workitem = index.cmp eq, %workitem, %c0 : index + scf.for %visibility_tile = [%c0 to %c4 step %c1] unroll { + %lane_outside_tile = index.cmp ne, %lane_key_tile, %visibility_tile : index + %lane_tile_visible = scalar.ori %lane_outside_tile, %lane_key_visible : i1 + %tile_value = scf.select %lane_tile_visible, %one_flag, %zero_flag : f32 + %subgroup_tile_visible = kernel.subgroup.reduce %tile_value : f32 + scf.if %is_first_workitem { + view.store %subgroup_tile_visible, %tile_visibility[%visibility_tile] : f32, view<4xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 @@ -38,24 +83,68 @@ template.def<@ggml.flash_attention.pv_block_wmma.prefill_full> device @ggml_flas %global_output_tile = index.assume %global_output_tile0 [lt(%global_output_tile0, %output_tile_count)] : index %output_tile_channel = index.mul %global_output_tile, %c64 : index %value_channel0 = index.add %value_head_base, %output_tile_channel : index - %value_channel = index.add %value_channel0, %subgroup_product_channel : index - %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> + %product_init = vector.fragment %c0_f32x4 shape [%m, %n] : vector<4xf32> + %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<4xf32>) -> (vector<4xf32>) unroll schedule(recurrence) { + %key_tile_ordinal = index.div %key_tile, %c16 : index + %tile_all_visible = scf.if %direct_value_layout -> (i1) { + %tile_visibility_value = view.load %tile_visibility[%key_tile_ordinal] : view<4xf32> -> f32 + %tile_has_no_hidden_keys = scalar.cmpf ogt, %tile_visibility_value, %zero_flag : f32 + scf.yield %tile_has_no_hidden_keys : i1 + } else { + %not_direct = scalar.constant false : i1 + scf.yield %not_direct : i1 + } + %next_product_accumulator = scf.if %tile_all_visible -> (vector<4xf32>) { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index + %value_channel = index.add %value_channel0, %subgroup_product_channel : index + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> vector<16xf16> + %direct_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %direct_product_accumulator : vector<4xf32> + } else { + scf.for %load_iteration = [%c0 to %c4 step %c1] { + %linear = index.madd %load_iteration, %c256, %workitem : index + %staged_row = index.div %linear, %c64 : index + %staged_channel = index.rem %linear, %c64 : index + %block_key0 = index.add %key_tile, %staged_row : index + %block_key = index.assume %block_key0 [range(%block_key0, 0, 63)] : index + %key_probabilities = vector.load %visibility[%block_key, %c0] : view<64x4xf16> -> vector<4xf16> + %key_probability = vector.reduce %key_probabilities, %zero_value : vector<4xf16>, f16 + %key_visible = scalar.cmpf ogt, %key_probability, %zero_value : f16 + %staged_value = scf.if %key_visible -> (f16) { + %source_token0 = index.add %key_origin, %block_key : index + %source_token = index.assume %source_token0 [lt(%source_token0, %bounded_key_value_token_count)] : index + %source_channel0 = index.add %value_channel0, %staged_channel : index + %source_channel = index.assume %source_channel0 [lt(%source_channel0, %value_width)] : index + %loaded = view.load %value_view[%source_token, %source_channel] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %zero_value : f16 + } + view.store %staged_value, %staged_values[%staged_row, %staged_channel] : f16, view<16x64xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %staged_values[%c0, %subgroup_product_channel] shape [%k, %n] : view<16x64xf16> -> vector<16xf16> + %staged_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %staged_product_accumulator : vector<4xf32> + } + scf.yield %next_product_accumulator : vector<4xf32> } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<4xf32>, view<16x64xf32> kernel.barrier scope(workgroup) ordering(acq_rel) %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output0_f32 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output0 = vector.fptrunc %block_output0_f32 : vector<4xf32> to vector<4xf16> + %block_output1_f32 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output1 = vector.fptrunc %block_output1_f32 : vector<4xf32> to vector<4xf16> + %block_output2_f32 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output2 = vector.fptrunc %block_output2_f32 : vector<4xf32> to vector<4xf16> + %block_output3_f32 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output3 = vector.fptrunc %block_output3_f32 : vector<4xf32> to vector<4xf16> %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> @@ -83,14 +172,33 @@ template.def<@ggml.flash_attention.pv_block_wmma.prefill_tail> device @ggml_flas %c32 = index.constant 32 : index %c64 = index.constant 64 : index %c0_offset = index.constant 0 : offset - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index %probability_layout = encoding.layout.strided [72, 1] : encoding %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x[%tail_stage_head_size]xf16> %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf32> + // Reuse score scratch after probability publication. + // Zero V when every query has zero probability to avoid signed-zero WMMA rounding. + %visibility_offset = index.constant 4096 : offset + %staged_value_offset = index.constant 2048 : offset + %visibility = buffer.view %product_stage[%visibility_offset] : buffer -> view<64x4xf16> + %staged_values = buffer.view %product_stage[%staged_value_offset] : buffer -> view<16x64xf16> + %zero_value = scalar.constant 0.0 : f16 + %workitem = kernel.workitem.id : index + %lane = kernel.subgroup.lane.id : index + %subgroup = kernel.subgroup.id : index + %c256 = index.constant 256 : index + %p0 = view.load %probability_stage_view[%query_row0, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p1 = view.load %probability_stage_view[%query_row1, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p2 = view.load %probability_stage_view[%query_row2, %lane] : view<16x64xf16, %probability_layout> -> f16 + %p3 = view.load %probability_stage_view[%query_row3, %lane] : view<16x64xf16, %probability_layout> -> f16 + %probabilities = vector.from_elements %p0, %p1, %p2, %p3 : vector<4xf16> + %maximum_probability = vector.reduce %probabilities, %zero_value : vector<4xf16>, f16 + view.store %maximum_probability, %visibility[%lane, %subgroup] : f16, view<64x4xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 @@ -108,22 +216,46 @@ template.def<@ggml.flash_attention.pv_block_wmma.prefill_tail> device @ggml_flas %output_tile_valid = index.cmp ult, %output_tile, %group_output_tile_count : index %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %output_tile_valid -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index - %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { + %tail_product_init = vector.fragment %c0_f32x4 shape [%m, %n] : vector<4xf32> + %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<4xf32>) -> (vector<4xf32>) unroll { + scf.for %load_iteration = [%c0 to %c4 step %c1] { + %linear = index.madd %load_iteration, %c256, %workitem : index + %staged_row = index.div %linear, %c64 : index + %staged_channel = index.rem %linear, %c64 : index + %block_key0 = index.add %key_tile, %staged_row : index + %block_key = index.assume %block_key0 [range(%block_key0, 0, 63)] : index + %key_probabilities = vector.load %visibility[%block_key, %c0] : view<64x4xf16> -> vector<4xf16> + %key_probability = vector.reduce %key_probabilities, %zero_value : vector<4xf16>, f16 + %key_visible = scalar.cmpf ogt, %key_probability, %zero_value : f16 + %staged_value = scf.if %key_visible -> (f16) { + %tail_key = index.assume %block_key [range(%block_key, 0, 31)] : index + %source_channel = index.add %output_tile_channel, %staged_channel : index + %loaded = view.load %tail_key_value_stage_view[%tail_key, %source_channel] : view<32x[%tail_stage_head_size]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %zero_value : f16 + } + view.store %staged_value, %staged_values[%staged_row, %staged_channel] : f16, view<16x64xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x[%tail_stage_head_size]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> + %value_fragment = vector.fragment.load %staged_values[%c0, %subgroup_product_channel] shape [%k, %n] : view<16x64xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_product_accumulator : vector<4xf32> } - vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<4xf32>, view<16x64xf32> kernel.barrier scope(workgroup) ordering(acq_rel) %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output0_f32 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output0 = vector.fptrunc %block_output0_f32 : vector<4xf32> to vector<4xf16> + %block_output1_f32 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output1 = vector.fptrunc %block_output1_f32 : vector<4xf32> to vector<4xf16> + %block_output2_f32 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output2 = vector.fptrunc %block_output2_f32 : vector<4xf32> to vector<4xf16> + %block_output3_f32 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf32> -> vector<4xf32> + %block_output3 = vector.fptrunc %block_output3_f32 : vector<4xf32> to vector<4xf16> %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom index e59762e94c0f..044a81add678 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom @@ -10,23 +10,25 @@ config.decl @ggml.flash_attention.attention_scale : f32 config.def @ggml.flash_attention.value_layout = 0 : index -template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) +config.def @ggml.flash_attention.head_batch_size = 1 : index -template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) +template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer, %query_offset: offset) + +template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer, %key_offset: offset) template.decl @ggml.flash_attention.qk_block_wmma.prefill_tail(%qk_head_size: index, %query_stage_stride: index, %tail_stage_head_size: index, %tail_score_wave: i1, %subgroup: index, %tail_key_value_stage: buffer, %query_stage: buffer, %score_stage: buffer) template.decl @ggml.flash_attention.online_softmax.prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) -template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer, %value_offset: offset) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) -template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer, %output_offset: offset, %f16_output_offset: offset) -template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %initial_max: f32, %initial_sum: f32, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer, %query_tile: index, %query_head0: index, %output_group0: index, %query_offset: offset, %key_offset: offset, %value_offset: offset, %mask_offset: offset, %output_offset: offset, %f16_output_offset: offset) -template.def<@ggml.flash_attention.prefill.stage_query> device @ggml_flash_attention_prefill_stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) { +template.def<@ggml.flash_attention.prefill.stage_query> device @ggml_flash_attention_prefill_stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer, %query_offset: offset) { %query_head_count = config.get @ggml.flash_attention.query_head_count : index %workitem = kernel.workitem.id : index %c0 = index.constant 0 : index @@ -34,7 +36,7 @@ template.def<@ggml.flash_attention.prefill.stage_query> device @ggml_flash_atten %c256 = index.constant 256 : index %c0_offset = index.constant 0 : offset %c0_f16 = scalar.constant 0.0 : f16 - %query_view = buffer.view %query[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%qk_head_size]xf32> + %query_view = buffer.view %query[%query_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%qk_head_size]xf32> %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x[%query_stage_stride]xf16> // Scale and truncate Q exactly once. The Vulkan reference does this before // entering its KV loop, making QK a native F16 WMMA while retaining F32 @@ -58,7 +60,7 @@ template.def<@ggml.flash_attention.prefill.stage_query> device @ggml_flash_atten template.return } -template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flash_attention_qk_block_wmma_prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) { +template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flash_attention_qk_block_wmma_prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer, %key_offset: offset) { %c0 = index.constant 0 : index %c16 = index.constant 16 : index %c0_offset = index.constant 0 : offset @@ -66,7 +68,7 @@ template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flas %n = index.constant 16 : index %k = index.constant 16 : index %query_transposed_layout = encoding.layout.strided [1, %query_stage_stride] : encoding - %key_view = buffer.view %key[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> + %key_view = buffer.view %key[%key_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x20xf32> %score_init_values = vector.constant 0.0 : vector<4xf32> @@ -83,8 +85,9 @@ template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flas template.return } -template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_prefill_body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) { - %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048)] : index +template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_prefill_body(%publish_f16: i1, %initial_max: f32, %initial_sum: f32, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer, %query_tile: index, %query_head0: index, %output_group0: index, %query_offset: offset, %key_offset: offset, %value_offset: offset, %mask_offset: offset, %output_offset: offset, %f16_output_offset: offset) { + %head_batch_size = config.get @ggml.flash_attention.head_batch_size : index + %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048), multiple_of(%query_token_count, %head_batch_size)] : index %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index %query_head_count = config.get @ggml.flash_attention.query_head_count : index %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index @@ -92,9 +95,6 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %qk_head_size = index.assume %qk_head_size0 [range(%qk_head_size0, 16, 512), multiple_of(%qk_head_size0, 16)] : index %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), multiple_of(%value_head_size0, 64)] : index - %query_tile = kernel.workgroup.id : index - %query_head0 = kernel.workgroup.id : index - %output_group0 = kernel.workgroup.id : index %query_head, %launch_query_head_count = index.assume %query_head0, %query_head_count [lt(%query_head0, %query_head_count)] : index, index %workitem = kernel.workitem.id : index %subgroup0 = kernel.subgroup.id : index @@ -126,6 +126,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %c0_f16x8 = vector.constant 0.0 : vector<8xf16> %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> + %initial_sum_vector = vector.splat %initial_sum : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index @@ -178,6 +179,36 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %query_valid1 = index.cmp ult, %query_token1, %bounded_query_token_count : index %query_valid2 = index.cmp ult, %query_token2, %bounded_query_token_count : index %query_valid3 = index.cmp ult, %query_token3, %bounded_query_token_count : index + %heads_as_queries = index.cmp ne, %head_batch_size, %c1 : index + %initial_max_vector = scf.if %heads_as_queries -> (vector<4xf32>) { + %sink_count = index.mul %head_batch_size, %query_head_count : index + %sinks = buffer.view %gate[%c0_offset] : buffer -> view<[%sink_count]xf32> + %head_in_batch0 = index.rem %query_token0, %head_batch_size : index + %i0 = index.madd %head_in_batch0, %query_head_count, %query_head : index + %head_in_batch1 = index.rem %query_token1, %head_batch_size : index + %i1 = index.madd %head_in_batch1, %query_head_count, %query_head : index + %head_in_batch2 = index.rem %query_token2, %head_batch_size : index + %i2 = index.madd %head_in_batch2, %query_head_count, %query_head : index + %head_in_batch3 = index.rem %query_token3, %head_batch_size : index + %i3 = index.madd %head_in_batch3, %query_head_count, %query_head : index + %s0 = scf.select %query_valid0, %i0, %c0 : index + %s1 = scf.select %query_valid1, %i1, %c0 : index + %s2 = scf.select %query_valid2, %i2, %c0 : index + %s3 = scf.select %query_valid3, %i3, %c0 : index + %b0 = index.assume %s0 [lt(%s0, %sink_count)] : index + %b1 = index.assume %s1 [lt(%s1, %sink_count)] : index + %b2 = index.assume %s2 [lt(%s2, %sink_count)] : index + %b3 = index.assume %s3 [lt(%s3, %sink_count)] : index + %v0 = view.load %sinks[%b0] : view<[%sink_count]xf32> -> f32 + %v1 = view.load %sinks[%b1] : view<[%sink_count]xf32> -> f32 + %v2 = view.load %sinks[%b2] : view<[%sink_count]xf32> -> f32 + %v3 = view.load %sinks[%b3] : view<[%sink_count]xf32> -> f32 + %values = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + scf.yield %values : vector<4xf32> + } else { + %values = vector.splat %initial_max : vector<4xf32> + scf.yield %values : vector<4xf32> + } %query_valid = vector.from_elements %query_valid0, %query_valid1, %query_valid2, %query_valid3 : vector<4xi1> %subgroup_product_channel = index.mul %subgroup, %c16 : index %lane_output_tile = index.div %lane, %c16 : index @@ -199,14 +230,19 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> + %key_view = buffer.view %key_aligned[%key_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> %value_layout_tag = config.get @ggml.flash_attention.value_layout : index %transposed_value = index.cmp eq, %value_layout_tag, %c1 : index %value_row_stride = scf.select %transposed_value, %c1, %value_width : index %value_column_stride = scf.select %transposed_value, %bounded_key_value_token_count, %c1 : index %value_layout = encoding.layout.strided [%value_row_stride, %value_column_stride] : encoding - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> + %value_view = buffer.view %value_aligned[%value_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> + %mask_query_count = index.div %bounded_query_token_count, %head_batch_size : index + %mask_query0 = index.div %query_token0, %head_batch_size : index + %mask_query1 = index.div %query_token1, %head_batch_size : index + %mask_query2 = index.div %query_token2, %head_batch_size : index + %mask_query3 = index.div %query_token3, %head_batch_size : index + %mask_view = buffer.view %mask_aligned[%mask_offset] : buffer -> view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> %query_stage = buffer.alloca align(16) %query_stage_bytes : buffer %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer @@ -219,9 +255,9 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr // before the next mask/score phase, so the two exchanges share storage. %product_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<16x64xf16> %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x[%tail_stage_head_size]xf16> - template.apply<@ggml.flash_attention.prefill.stage_query>(%bounded_query_token_count, %query_origin, %query_head, %query_load_iteration_count, %query_stage_stride, %qk_head_size, %attention_scale, %query_aligned, %query_stage) : (index, index, index, index, index, index, f32, buffer, buffer) + template.apply<@ggml.flash_attention.prefill.stage_query>(%bounded_query_token_count, %query_origin, %query_head, %query_load_iteration_count, %query_stage_stride, %qk_head_size, %attention_scale, %query_aligned, %query_stage, %query_offset) : (index, index, index, index, index, index, f32, buffer, buffer, offset) kernel.barrier scope(workgroup) ordering(acq_rel) - %full_max, %full_sum, %full_output0, %full_output1, %full_output2, %full_output3 = scf.for %key_origin = [%c0 to %full_key_value_token_count step %c64](%current_max = %negative_f32x4 : vector<4xf32>, %current_sum = %c0_f32x4 : vector<4xf32>, %current_output0 = %output_zero0 : vector<4xf16>, %current_output1 = %output_zero1 : vector<4xf16>, %current_output2 = %output_zero2 : vector<4xf16>, %current_output3 = %output_zero3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %full_max, %full_sum, %full_output0, %full_output1, %full_output2, %full_output3 = scf.for %key_origin = [%c0 to %full_key_value_token_count step %c64](%current_max = %initial_max_vector : vector<4xf32>, %current_sum = %initial_sum_vector : vector<4xf32>, %current_output0 = %output_zero0 : vector<4xf16>, %current_output1 = %output_zero1 : vector<4xf16>, %current_output2 = %output_zero2 : vector<4xf16>, %current_output3 = %output_zero3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { // Cache the complete 16x64 mask tile before QK. Causal masks leave future // KV blocks entirely at negative infinity; reducing the cached tile lets // the workgroup skip QK, softmax, and P*V for those blocks. This mirrors @@ -229,25 +265,25 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %key_token0 = index.add %key_origin, %lane : index %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index %mask0 = scf.if %query_valid0 -> (f16) { - %value = view.load %mask_view[%query_token0, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%mask_query0, %key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 scf.yield %value : f16 } else { scf.yield %c0_f16 : f16 } %mask1 = scf.if %query_valid1 -> (f16) { - %value = view.load %mask_view[%query_token1, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%mask_query1, %key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 scf.yield %value : f16 } else { scf.yield %c0_f16 : f16 } %mask2 = scf.if %query_valid2 -> (f16) { - %value = view.load %mask_view[%query_token2, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%mask_query2, %key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 scf.yield %value : f16 } else { scf.yield %c0_f16 : f16 } %mask3 = scf.if %query_valid3 -> (f16) { - %value = view.load %mask_view[%query_token3, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%mask_query3, %key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 scf.yield %value : f16 } else { scf.yield %c0_f16 : f16 @@ -263,6 +299,8 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr } kernel.barrier scope(workgroup) ordering(acq_rel) %subgroup_mask_maxima = vector.load %mask_summary_view[%c0] : view<4xf32> -> vector<4xf32> + // Retire summary reads before scores or the next masked block reuse LDS. + kernel.barrier scope(workgroup) ordering(acq_rel) %workgroup_mask_maximum = vector.reduce %subgroup_mask_maxima, %negative_large : vector<4xf32>, f32 %block_has_attention = scalar.cmpf ogt, %workgroup_mask_maximum, %negative_large : f32 %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 = scf.if %block_has_attention -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { @@ -270,7 +308,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index - template.apply<@ggml.flash_attention.qk_block_wmma.prefill_full>(%bounded_key_value_token_count, %key_width, %qk_head_size, %query_stage_stride, %score_key_origin, %key_head_base, %subgroup_score_column, %key_aligned, %query_stage, %score_stage) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + template.apply<@ggml.flash_attention.qk_block_wmma.prefill_full>(%bounded_key_value_token_count, %key_width, %qk_head_size, %query_stage_stride, %score_key_origin, %key_head_base, %subgroup_score_column, %key_aligned, %query_stage, %score_stage, %key_offset) : (index, index, index, index, index, index, index, buffer, buffer, buffer, offset) // LDS transposes ownership from one 16-column score slice per wave to // four complete query rows per wave. Every lane contributes one key // column to each of those rows. @@ -286,7 +324,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr // Compute the 64-channel tiles in this output group sequentially. This // keeps one P*V accumulator live per wave and reuses a 2 KiB exchange // tile instead of retaining multiple output groups. - %next_output0, %next_output1, %next_output2, %next_output3 = template.apply<@ggml.flash_attention.pv_block_wmma.prefill_full>(%old_scale, %current_output0, %current_output1, %current_output2, %current_output3, %group_output_tile_count, %output_group_tile_base, %output_tile_count, %value_head_base, %subgroup_product_channel, %lane_output_tile, %lane_product_channel, %query_row0, %query_row1, %query_row2, %query_row3, %key_origin, %last_full_key_tile_start, %bounded_key_value_token_count, %value_width, %value_row_stride, %value_column_stride, %value_aligned, %probability_stage, %score_stage) : (vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + %next_output0, %next_output1, %next_output2, %next_output3 = template.apply<@ggml.flash_attention.pv_block_wmma.prefill_full>(%old_scale, %current_output0, %current_output1, %current_output2, %current_output3, %group_output_tile_count, %output_group_tile_base, %output_tile_count, %value_head_base, %subgroup_product_channel, %lane_output_tile, %lane_product_channel, %query_row0, %query_row1, %query_row2, %query_row3, %key_origin, %last_full_key_tile_start, %bounded_key_value_token_count, %value_width, %value_row_stride, %value_column_stride, %value_aligned, %probability_stage, %score_stage, %value_offset) : (vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, offset) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> } else { scf.yield %current_max, %current_sum, %current_output0, %current_output1, %current_output2, %current_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> @@ -333,7 +371,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %masked_score0 = scf.if %tail_valid0 -> (f32) { %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x20xf32> -> f32 - %mask_f16 = view.load %mask_view[%query_token0, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f16 = view.load %mask_view[%mask_query0, %tail_key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 %mask_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -343,7 +381,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %masked_score1 = scf.if %tail_valid1 -> (f32) { %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x20xf32> -> f32 - %mask_f16 = view.load %mask_view[%query_token1, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f16 = view.load %mask_view[%mask_query1, %tail_key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 %mask_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -353,7 +391,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %masked_score2 = scf.if %tail_valid2 -> (f32) { %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x20xf32> -> f32 - %mask_f16 = view.load %mask_view[%query_token2, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f16 = view.load %mask_view[%mask_query2, %tail_key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 %mask_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -363,7 +401,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %masked_score3 = scf.if %tail_valid3 -> (f32) { %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x20xf32> -> f32 - %mask_f16 = view.load %mask_view[%query_token3, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f16 = view.load %mask_view[%mask_query3, %tail_key_token] : view<[%mask_query_count]x[%bounded_key_value_token_count]xf16> -> f16 %mask_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -402,7 +440,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %next_output0, %next_output1, %next_output2, %next_output3 = template.apply<@ggml.flash_attention.pv_block_wmma.prefill_tail>(%old_scale, %current_output0, %current_output1, %current_output2, %current_output3, %group_output_tile_count, %subgroup_product_channel, %lane_output_tile, %lane_product_channel, %query_row0, %query_row1, %query_row2, %query_row3, %tail_stage_head_size, %tail_key_value_stage, %probability_stage, %score_stage) : (vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> } - template.apply<@ggml.flash_attention.prefill.publish_output>(%publish_f16, %bounded_query_token_count, %query_head_count, %value_head_size, %lane_has_output, %query_valid0, %query_valid1, %query_valid2, %query_valid3, %query_token0, %query_token1, %query_token2, %query_token3, %query_head, %lane_output_channel, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3, %gate, %output_aligned, %f16_output) : (i1, index, index, index, i1, i1, i1, i1, i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, buffer, buffer, buffer) + template.apply<@ggml.flash_attention.prefill.publish_output>(%publish_f16, %bounded_query_token_count, %query_head_count, %value_head_size, %lane_has_output, %query_valid0, %query_valid1, %query_valid2, %query_valid3, %query_token0, %query_token1, %query_token2, %query_token3, %query_head, %lane_output_channel, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3, %gate, %output_aligned, %f16_output, %output_offset, %f16_output_offset) : (i1, index, index, index, i1, i1, i1, i1, i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, buffer, buffer, buffer, offset, offset) template.return } 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 index 4b4eb1c92e95..09390b3b6f2c 100644 --- 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 @@ -62,11 +62,11 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %f32_element_bytes = index.constant 4 : 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 + %wave_result_stage_bytes = index.constant 1024 : offset + %result_stage_bytes = index.constant 2048 : 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> + %zero_accumulator = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index @@ -95,8 +95,8 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %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> + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32> %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 @@ -105,17 +105,17 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %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> + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> %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>) %reuse_q5_k = index.constant 5 : index %reuse_iq4_xs = index.constant 23 : index %weight_is_q5 = index.cmp eq, %weight_format, %reuse_q5_k : index %weight_is_iq4 = index.cmp eq, %weight_format, %reuse_iq4_xs : index - %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>) { + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<4xf32>, %block_acc01 = %init01 : vector<4xf32>, %block_acc10 = %init10 : vector<4xf32>, %block_acc11 = %init11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { // Retain weight Q5_K headers and packet high bits for all eight K32 groups; payloads are still loaded per group. %weight_dm, %weight_scales, %weight_qh = scf.if %weight_is_q5 -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { %zero_dm = vector.constant 0.0 : vector<4x2xf16> @@ -174,7 +174,7 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %empty_scales_l = vector.constant 0 : vector<4xi32> scf.yield %empty_header, %empty_scales_l : vector<4xi32>, vector<4xi32> } - %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_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<4xf32>, %acc01 = %block_acc01 : vector<4xf32>, %acc10 = %block_acc10 : vector<4xf32>, %acc11 = %block_acc11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { %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 @@ -247,21 +247,21 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ } } 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 { + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<4xf32>, %half_acc01 = %acc01 : vector<4xf32>, %half_acc10 = %acc10 : vector<4xf32>, %half_acc11 = %acc11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) 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> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %next00, %next01, %next10, %next11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } %publish_token0 = index.div %lane, %c4 : index %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index @@ -285,35 +285,31 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %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> + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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> + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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> + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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> + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom index 0ffae0f3d5e2..db615c5ee212 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom @@ -57,11 +57,11 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%ful %f32_element_bytes = index.constant 4 : offset %weight_stage_bytes = index.constant 5120 : offset %activation_stage_bytes = index.constant 5120 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 2048 : offset + %wave_result_stage_bytes = index.constant 1024 : offset + %result_stage_bytes = index.constant 4096 : 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> + %zero_accumulator = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index @@ -90,8 +90,8 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%ful %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x64xf16, %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> + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32> %channel_tile_base = index.mul %channel_tile, %c64 : index %token_tile_base = index.mul %token_tile, %c64 : index %load_packet = index.rem %workitem, %c8 : index @@ -99,14 +99,14 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%ful %load_row0 = index.div %workitem, %c8 : index %load_row = index.assume %load_row0 [range(%load_row0, 0, 31)] : index %subgroup_channel_add = index.mul %subgroup, %c16 : index - %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init02 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init03 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init02 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init03 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> %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, %result02, %result03 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc02 = %init02 : vector<8xf16>, %block_acc03 = %init03 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %block_result00, %block_result01, %block_result02, %block_result03 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc02 = %block_acc02 : vector<8xf16>, %acc03 = %block_acc03 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %result00, %result01, %result02, %result03 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<4xf32>, %block_acc01 = %init01 : vector<4xf32>, %block_acc02 = %init02 : vector<4xf32>, %block_acc03 = %init03 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block_result00, %block_result01, %block_result02, %block_result03 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<4xf32>, %acc01 = %block_acc01 : vector<4xf32>, %acc02 = %block_acc02 : vector<4xf32>, %acc03 = %block_acc03 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { %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 @@ -155,22 +155,22 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%ful vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<64x40xf16> } kernel.barrier scope(workgroup) ordering(acq_rel) - %next00, %next01, %next02, %next03 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc02 = %acc02 : vector<8xf16>, %half_acc03 = %acc03 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %next00, %next01, %next02, %next03 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<4xf32>, %half_acc01 = %acc01 : vector<4xf32>, %half_acc02 = %acc02 : vector<4xf32>, %half_acc03 = %acc03 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) unroll { %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x64xf16, %activation_fragment_layout> -> vector<16xf16> %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x64xf16, %activation_fragment_layout> -> vector<16xf16> %rhs2 = vector.fragment.load %activation_fragment_view[%k_half, %c32] shape [%k, %n] : view<32x64xf16, %activation_fragment_layout> -> vector<16xf16> %rhs3 = vector.fragment.load %activation_fragment_view[%k_half, %c48] shape [%k, %n] : view<32x64xf16, %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_next02 = vector.mma %lhs0, %rhs2, %half_acc02 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next03 = vector.mma %lhs0, %rhs3, %half_acc03 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %half_next00, %half_next01, %half_next02, %half_next03 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next02 = vector.mma %lhs0, %rhs2, %half_acc02 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next03 = vector.mma %lhs0, %rhs3, %half_acc03 : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %half_next00, %half_next01, %half_next02, %half_next03 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next02, %next03 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %next00, %next01, %next02, %next03 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } - scf.yield %block_result00, %block_result01, %block_result02, %block_result03 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %block_result00, %block_result01, %block_result02, %block_result03 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } %publish_token0 = index.div %lane, %c4 : index %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index @@ -199,35 +199,31 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%ful %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 %writes02 = scalar.andi %valid_token2, %valid_channel0 : i1 %writes03 = scalar.andi %valid_token3, %valid_channel0 : i1 - vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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> + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> 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 %result02, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %result02, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes02 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes02, %token2, %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 %result03, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %result03, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %result_fragment_layout> kernel.barrier scope(subgroup) ordering(acq_rel) scf.if %writes03 { - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes03, %token3, %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) } 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) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/mxfp4_prefill_prefetch.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/mxfp4_prefill_prefetch.loom new file mode 100644 index 000000000000..eca627de5206 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/mxfp4_prefill_prefetch.loom @@ -0,0 +1,550 @@ +func.def inline @ggml_mxfp4_prefetch_panel(%weight_noalias: buffer, %input_noalias: buffer, %route_stage: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %channel_tile_base: index, %bounded_output_size: index, %quant_groups: index, %load_row: index, %load_k: index, %weight_klo: index, %bounded_route_count: index, %bounded_input_route_count: index, %block_base: index) -> (i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32) { + %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 + %c32 = index.constant 32 : index + %zero_offset = index.constant 0 : offset + %one_byte = index.constant 1 : offset + %c17_bytes = index.constant 17 : offset + %header_bytes = index.constant 16 : offset + %group_bytes = index.constant 144 : offset + %zi8 = scalar.constant 0 : i32 + %zh = scalar.constant 0 : i32 + %zi32 = scalar.constant 0 : i32 + %z4 = scalar.constant 0 : i32 + %z8 = vector.constant 0 : vector<2xi32> + %route_view = buffer.view %route_stage[%zero_offset] : buffer -> view<32xi32> + %zero_wp = vector.constant 0 : vector<2x4xi32> + %zero_we = vector.constant 0 : vector<2x4xi32> + %zero_ap = vector.constant 0 : vector<2x2x2xi32> + %zero_ad = vector.constant 0 : vector<2x2xi32> + %out_wp, %out_we, %out_ap, %out_ad = scf.for %ks = [%c0 to %c2 step %c1](%bank_wp = %zero_wp : vector<2x4xi32>, %bank_we = %zero_we : vector<2x4xi32>, %bank_ap = %zero_ap : vector<2x2x2xi32>, %bank_ad = %zero_ad : vector<2x2xi32>) -> (vector<2x4xi32>, vector<2x4xi32>, vector<2x2x2xi32>, vector<2x2xi32>) unroll { + %block = index.add %block_base, %ks : index + %bo = index.scale %block, %c17_bytes : index, offset -> offset + %next_wp, %next_we = scf.for %slot = [%c0 to %c4 step %c1](%pw = %bank_wp : vector<2x4xi32>, %pe = %bank_we : vector<2x4xi32>) -> (vector<2x4xi32>, vector<2x4xi32>) unroll { + %row = index.madd %slot, %c16, %load_row : index + %channel = index.add %channel_tile_base, %row : index + %valid = index.cmp ult, %channel, %bounded_output_size : index + %co = index.scale %channel, %weight_row_bytes : index, offset -> offset + %ro = index.add %expert_byte_base, %co : offset + %wb = index.add %ro, %bo : offset + %pb = index.add %wb, %one_byte : offset + %sv = buffer.view %weight_noalias[%wb] : buffer -> view<4xi8> + %qv = buffer.view %weight_noalias[%pb] : buffer -> view<16xi8> + %eb = vector.load %sv[%c0] : view<4xi8> -> vector<4xi8> + %ew = vector.bitcast %eb : vector<4xi8> to vector<1xi32> + %e = vector.extract %ew[%c0] : vector<1xi32> -> i32 + %qb = vector.load %qv[%weight_klo] : view<16xi8> -> vector<4xi8> + %qw = vector.bitcast %qb : vector<4xi8> to vector<1xi32> + %q = vector.extract %qw[%c0] : vector<1xi32> -> i32 + %payload = scalar.ori %q, %zi32 : i32 + %exponent = scalar.ori %e, %zi32 : i32 + %nw = vector.insert %payload into %pw[%ks, %slot] : i32, vector<2x4xi32> + %ne = vector.insert %exponent into %pe[%ks, %slot] : i32, vector<2x4xi32> + scf.yield %nw, %ne : vector<2x4xi32>, vector<2x4xi32> + } + %next_ap, %next_ad = scf.for %slot = [%c0 to %c2 step %c1](%pa = %bank_ap : vector<2x2x2xi32>, %pd = %bank_ad : vector<2x2xi32>) -> (vector<2x2x2xi32>, vector<2x2xi32>) unroll { + %row0 = index.madd %slot, %c16, %load_row : index + %row = index.assume %row0 [range(%row0, 0, 31)] : index + %assignment_i32 = view.load %route_view[%row] : view<32xi32> -> i32 + %valid = scalar.cmpi sge, %assignment_i32, %zi32 : i32 + %safe_assignment = scf.select %valid, %assignment_i32, %zi32 : i32 + %assignment = index.cast %safe_assignment : i32 to index + %token = index.div %assignment, %bounded_route_count : index + %route = index.rem %assignment, %bounded_route_count : index + %input_route = index.rem %route, %bounded_input_route_count : index + %input_row = index.madd %token, %bounded_input_route_count, %input_route : index + %groups_per_row = index.div %quant_groups, %c4 : index + %group = index.div %block, %c4 : index + %local_block = index.rem %block, %c4 : index + %global_group = index.madd %input_row, %groups_per_row, %group : index + %group_offset = index.scale %global_group, %group_bytes : index, offset -> offset + %pb = index.add %group_offset, %header_bytes : offset + %payload_i = index.madd %local_block, %c32, %load_k : index + %ds_i = index.mul %local_block, %c2 : index + %iv = buffer.view %input_noalias[%pb] : buffer -> view<128xi8> + %dv = buffer.view %input_noalias[%group_offset] : buffer -> view<8xf16> + %vb = vector.load %iv[%payload_i] : view<128xi8> -> vector<8xi8> + %v = vector.bitcast %vb : vector<8xi8> to vector<2xi32> + %dh = vector.load %dv[%ds_i] : view<8xf16> -> vector<2xf16> + %dw = vector.bitcast %dh : vector<2xf16> to vector<1xi32> + %d = vector.extract %dw[%c0] : vector<1xi32> -> i32 + %payload = vector.bitcast %v : vector<2xi32> to vector<2xi32> + %scale = scalar.ori %d, %zi32 : i32 + %na = vector.insert %payload into %pa[%ks, %slot] : vector<2xi32>, vector<2x2x2xi32> + %nd = vector.insert %scale into %pd[%ks, %slot] : i32, vector<2x2xi32> + scf.yield %na, %nd : vector<2x2x2xi32>, vector<2x2xi32> + } + scf.yield %next_wp, %next_we, %next_ap, %next_ad : vector<2x4xi32>, vector<2x4xi32>, vector<2x2x2xi32>, vector<2x2xi32> + } + %c3 = index.constant 3 : index + %ret_wp00 = vector.extract %out_wp[%c0, %c0] : vector<2x4xi32> -> i32 + %ret_wp01 = vector.extract %out_wp[%c0, %c1] : vector<2x4xi32> -> i32 + %ret_wp02 = vector.extract %out_wp[%c0, %c2] : vector<2x4xi32> -> i32 + %ret_wp03 = vector.extract %out_wp[%c0, %c3] : vector<2x4xi32> -> i32 + %ret_wp10 = vector.extract %out_wp[%c1, %c0] : vector<2x4xi32> -> i32 + %ret_wp11 = vector.extract %out_wp[%c1, %c1] : vector<2x4xi32> -> i32 + %ret_wp12 = vector.extract %out_wp[%c1, %c2] : vector<2x4xi32> -> i32 + %ret_wp13 = vector.extract %out_wp[%c1, %c3] : vector<2x4xi32> -> i32 + %ret_we00 = vector.extract %out_we[%c0, %c0] : vector<2x4xi32> -> i32 + %ret_we01 = vector.extract %out_we[%c0, %c1] : vector<2x4xi32> -> i32 + %ret_we02 = vector.extract %out_we[%c0, %c2] : vector<2x4xi32> -> i32 + %ret_we03 = vector.extract %out_we[%c0, %c3] : vector<2x4xi32> -> i32 + %ret_we10 = vector.extract %out_we[%c1, %c0] : vector<2x4xi32> -> i32 + %ret_we11 = vector.extract %out_we[%c1, %c1] : vector<2x4xi32> -> i32 + %ret_we12 = vector.extract %out_we[%c1, %c2] : vector<2x4xi32> -> i32 + %ret_we13 = vector.extract %out_we[%c1, %c3] : vector<2x4xi32> -> i32 + %ret_ap000 = vector.extract %out_ap[%c0, %c0, %c0] : vector<2x2x2xi32> -> i32 + %ret_ap001 = vector.extract %out_ap[%c0, %c0, %c1] : vector<2x2x2xi32> -> i32 + %ret_ap010 = vector.extract %out_ap[%c0, %c1, %c0] : vector<2x2x2xi32> -> i32 + %ret_ap011 = vector.extract %out_ap[%c0, %c1, %c1] : vector<2x2x2xi32> -> i32 + %ret_ap100 = vector.extract %out_ap[%c1, %c0, %c0] : vector<2x2x2xi32> -> i32 + %ret_ap101 = vector.extract %out_ap[%c1, %c0, %c1] : vector<2x2x2xi32> -> i32 + %ret_ap110 = vector.extract %out_ap[%c1, %c1, %c0] : vector<2x2x2xi32> -> i32 + %ret_ap111 = vector.extract %out_ap[%c1, %c1, %c1] : vector<2x2x2xi32> -> i32 + %ret_ad00 = vector.extract %out_ad[%c0, %c0] : vector<2x2xi32> -> i32 + %ret_ad01 = vector.extract %out_ad[%c0, %c1] : vector<2x2xi32> -> i32 + %ret_ad10 = vector.extract %out_ad[%c1, %c0] : vector<2x2xi32> -> i32 + %ret_ad11 = vector.extract %out_ad[%c1, %c1] : vector<2x2xi32> -> i32 + func.return %ret_wp00, %ret_wp01, %ret_wp02, %ret_wp03, %ret_wp10, %ret_wp11, %ret_wp12, %ret_wp13, %ret_we00, %ret_we01, %ret_we02, %ret_we03, %ret_we10, %ret_we11, %ret_we12, %ret_we13, %ret_ap000, %ret_ap001, %ret_ap010, %ret_ap011, %ret_ap100, %ret_ap101, %ret_ap110, %ret_ap111, %ret_ad00, %ret_ad01, %ret_ad10, %ret_ad11 : i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32 +} + +func.def inline @ggml_mxfp4_decode_packet(%p: vector<4xi8>, %scale_byte: i8) -> (vector<4xi8>, vector<4xi8>, f32) { + %exp = scalar.extui %scale_byte : i8 to i32 + %shift23 = scalar.constant 23 : i32 + %iz = scalar.constant 0 : i32 + %sub = scalar.constant 4194304 : i32 + %nan_exp = scalar.constant 255 : i32 + %nan_bits = scalar.constant 2143289344 : i32 + %bits = scalar.shli %exp, %shift23 : i32 + %is_zero = scalar.cmpi eq, %exp, %iz : i32 + %is_nan = scalar.cmpi eq, %exp, %nan_exp : i32 + %finite = scf.select %is_zero, %sub, %bits : i32 + %scale_bits = scf.select %is_nan, %nan_bits, %finite : i32 + %scale_raw = scalar.bitcast %scale_bits : i32 to f32 + %half = scalar.constant 0.5 : f32 + %scale = scalar.mulf %scale_raw, %half : f32 + %pw = vector.bitcast %p : vector<4xi8> to vector<1xi32> + %s4 = vector.constant 4 : vector<1xi32> + %mask = vector.constant 15 : vector<4xi8> + %hiw = vector.shrui %pw, %s4 : vector<1xi32> + %hib = vector.bitcast %hiw : vector<1xi32> to vector<4xi8> + %li = vector.andi %p, %mask : vector<4xi8> + %hi = vector.andi %hib, %mask : vector<4xi8> + %c0 = scalar.constant 0 : i8 + %c1 = scalar.constant 1 : i8 + %c2 = scalar.constant 2 : i8 + %c3 = scalar.constant 3 : i8 + %c4 = scalar.constant 4 : i8 + %c5 = scalar.constant 6 : i8 + %c6 = scalar.constant 8 : i8 + %c7 = scalar.constant 12 : i8 + %c8 = scalar.constant 0 : i8 + %c9 = scalar.constant -1 : i8 + %c10 = scalar.constant -2 : i8 + %c11 = scalar.constant -3 : i8 + %c12 = scalar.constant -4 : i8 + %c13 = scalar.constant -6 : i8 + %c14 = scalar.constant -8 : i8 + %c15 = scalar.constant -12 : i8 + %table = vector.from_elements %c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7, %c8, %c9, %c10, %c11, %c12, %c13, %c14, %c15 : vector<16xi8> + %lv = vector.table.lookup %table[%li] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + %hv = vector.table.lookup %table[%hi] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + func.return %lv, %hv, %scale : vector<4xi8>, vector<4xi8>, f32 +} + +// Shared routed expert WMMA matmul pipeline for GGML MUL_MAT_ID. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul_id.mxfp4_prefetch.publish_vector8 +// ggml.matmul_id.mxfp4_prefetch.finish_tile +template.decl @ggml.matmul_id.mxfp4_prefetch.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.mxfp4_prefetch.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.mxfp4_prefetch.publish_vector8(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<8xf32>, %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 + %c127_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, %c127_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.mxfp4_prefetch.core> device @ggml_mul_mat_id_mxfp4_prefetch_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 + %c80 = index.constant 80 : index + %c64 = index.constant 64 : index + %c63 = index.constant 63 : index + %c128 = index.constant 128 : 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 1024 : offset + %result_stage_bytes = index.constant 2048 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c2 = index.constant 2 : index + %c17_bytes = index.constant 17 : offset + %one_byte = index.constant 1 : offset + %header_bytes = index.constant 16 : offset + %group_bytes = index.constant 144 : offset + %scale_w_bytes = index.constant 512 : offset + %scale_a_bytes = index.constant 256 : offset + %weight_scales = buffer.alloca align(16) %scale_w_bytes : buffer + %input_scales = buffer.alloca align(16) %scale_a_bytes : buffer + %ws = buffer.view %weight_scales[%c0_offset] : buffer -> view<2x64xf32> + %xs = buffer.view %input_scales[%c0_offset] : buffer -> view<2x32xf32> + %zero_i8x4 = vector.constant 0 : vector<4xi8> + %zero_i8x8 = vector.constant 0 : vector<8xi8> + %zero_float = scalar.constant 0.0 : f32 + %zero_integer = vector.constant 0 : vector<8xi32> + %i8_schema = encoding.define #encoding.operand : encoding + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %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 = index.constant 136 : 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 + + %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<64x80xi8> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x80xi8> + %activation_fragment_layout = encoding.layout.strided [1, %c80] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x32xi8, %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<16x16xf32, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32> + %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, %c4 : index + %load_k = index.mul %load_packet, %c8 : index + %weight_klo = index.mul %load_packet, %c4 : index + %weight_khi = index.add %weight_klo, %c16 : index + %load_row0 = index.div %workitem, %c4 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %scale_writer = index.cmp eq, %load_packet, %c0 : 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 + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %lane_row_add = index.mul %lane_hi, %c8 : index + %quant_groups = index.div %input_size, %c32 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %bank_c3 = index.constant 3 : index + %initial_wp00, %initial_wp01, %initial_wp02, %initial_wp03, %initial_wp10, %initial_wp11, %initial_wp12, %initial_wp13, %initial_we00, %initial_we01, %initial_we02, %initial_we03, %initial_we10, %initial_we11, %initial_we12, %initial_we13, %initial_ap000, %initial_ap001, %initial_ap010, %initial_ap011, %initial_ap100, %initial_ap101, %initial_ap110, %initial_ap111, %initial_ad00, %initial_ad01, %initial_ad10, %initial_ad11 = func.call @ggml_mxfp4_prefetch_panel(%weight_noalias, %input_noalias, %route_stage, %expert_byte_base, %weight_row_bytes, %channel_tile_base, %bounded_output_size, %quant_groups, %load_row, %load_k, %weight_klo, %bounded_route_count, %bounded_input_route_count, %c0) : (buffer, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index) -> (i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32) + %result00, %result01, %result10, %result11, %unused_wp00, %unused_wp01, %unused_wp02, %unused_wp03, %unused_wp10, %unused_wp11, %unused_wp12, %unused_wp13, %unused_we00, %unused_we01, %unused_we02, %unused_we03, %unused_we10, %unused_we11, %unused_we12, %unused_we13, %unused_ap000, %unused_ap001, %unused_ap010, %unused_ap011, %unused_ap100, %unused_ap101, %unused_ap110, %unused_ap111, %unused_ad00, %unused_ad01, %unused_ad10, %unused_ad11 = scf.for %block_base = [%c0 to %quant_groups step %c2](%acc00 = %init : vector<8xf32>, %acc01 = %init : vector<8xf32>, %acc10 = %init : vector<8xf32>, %acc11 = %init : vector<8xf32>, %carry_wp00 = %initial_wp00 : i32, %carry_wp01 = %initial_wp01 : i32, %carry_wp02 = %initial_wp02 : i32, %carry_wp03 = %initial_wp03 : i32, %carry_wp10 = %initial_wp10 : i32, %carry_wp11 = %initial_wp11 : i32, %carry_wp12 = %initial_wp12 : i32, %carry_wp13 = %initial_wp13 : i32, %carry_we00 = %initial_we00 : i32, %carry_we01 = %initial_we01 : i32, %carry_we02 = %initial_we02 : i32, %carry_we03 = %initial_we03 : i32, %carry_we10 = %initial_we10 : i32, %carry_we11 = %initial_we11 : i32, %carry_we12 = %initial_we12 : i32, %carry_we13 = %initial_we13 : i32, %carry_ap000 = %initial_ap000 : i32, %carry_ap001 = %initial_ap001 : i32, %carry_ap010 = %initial_ap010 : i32, %carry_ap011 = %initial_ap011 : i32, %carry_ap100 = %initial_ap100 : i32, %carry_ap101 = %initial_ap101 : i32, %carry_ap110 = %initial_ap110 : i32, %carry_ap111 = %initial_ap111 : i32, %carry_ad00 = %initial_ad00 : i32, %carry_ad01 = %initial_ad01 : i32, %carry_ad10 = %initial_ad10 : i32, %carry_ad11 = %initial_ad11 : i32) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32) { + %zero_bank_wp = vector.constant 0 : vector<2x4xi32> + %built_wp00 = vector.insert %carry_wp00 into %zero_bank_wp[%c0, %c0] : i32, vector<2x4xi32> + %built_wp01 = vector.insert %carry_wp01 into %built_wp00[%c0, %c1] : i32, vector<2x4xi32> + %built_wp02 = vector.insert %carry_wp02 into %built_wp01[%c0, %c2] : i32, vector<2x4xi32> + %built_wp03 = vector.insert %carry_wp03 into %built_wp02[%c0, %bank_c3] : i32, vector<2x4xi32> + %built_wp10 = vector.insert %carry_wp10 into %built_wp03[%c1, %c0] : i32, vector<2x4xi32> + %built_wp11 = vector.insert %carry_wp11 into %built_wp10[%c1, %c1] : i32, vector<2x4xi32> + %built_wp12 = vector.insert %carry_wp12 into %built_wp11[%c1, %c2] : i32, vector<2x4xi32> + %bank_wp = vector.insert %carry_wp13 into %built_wp12[%c1, %bank_c3] : i32, vector<2x4xi32> + %zero_bank_we = vector.constant 0 : vector<2x4xi32> + %built_we00 = vector.insert %carry_we00 into %zero_bank_we[%c0, %c0] : i32, vector<2x4xi32> + %built_we01 = vector.insert %carry_we01 into %built_we00[%c0, %c1] : i32, vector<2x4xi32> + %built_we02 = vector.insert %carry_we02 into %built_we01[%c0, %c2] : i32, vector<2x4xi32> + %built_we03 = vector.insert %carry_we03 into %built_we02[%c0, %bank_c3] : i32, vector<2x4xi32> + %built_we10 = vector.insert %carry_we10 into %built_we03[%c1, %c0] : i32, vector<2x4xi32> + %built_we11 = vector.insert %carry_we11 into %built_we10[%c1, %c1] : i32, vector<2x4xi32> + %built_we12 = vector.insert %carry_we12 into %built_we11[%c1, %c2] : i32, vector<2x4xi32> + %bank_we = vector.insert %carry_we13 into %built_we12[%c1, %bank_c3] : i32, vector<2x4xi32> + %zero_bank_ap = vector.constant 0 : vector<2x2x2xi32> + %built_ap000 = vector.insert %carry_ap000 into %zero_bank_ap[%c0, %c0, %c0] : i32, vector<2x2x2xi32> + %built_ap001 = vector.insert %carry_ap001 into %built_ap000[%c0, %c0, %c1] : i32, vector<2x2x2xi32> + %built_ap010 = vector.insert %carry_ap010 into %built_ap001[%c0, %c1, %c0] : i32, vector<2x2x2xi32> + %built_ap011 = vector.insert %carry_ap011 into %built_ap010[%c0, %c1, %c1] : i32, vector<2x2x2xi32> + %built_ap100 = vector.insert %carry_ap100 into %built_ap011[%c1, %c0, %c0] : i32, vector<2x2x2xi32> + %built_ap101 = vector.insert %carry_ap101 into %built_ap100[%c1, %c0, %c1] : i32, vector<2x2x2xi32> + %built_ap110 = vector.insert %carry_ap110 into %built_ap101[%c1, %c1, %c0] : i32, vector<2x2x2xi32> + %bank_ap = vector.insert %carry_ap111 into %built_ap110[%c1, %c1, %c1] : i32, vector<2x2x2xi32> + %zero_bank_ad = vector.constant 0 : vector<2x2xi32> + %built_ad00 = vector.insert %carry_ad00 into %zero_bank_ad[%c0, %c0] : i32, vector<2x2xi32> + %built_ad01 = vector.insert %carry_ad01 into %built_ad00[%c0, %c1] : i32, vector<2x2xi32> + %built_ad10 = vector.insert %carry_ad10 into %built_ad01[%c1, %c0] : i32, vector<2x2xi32> + %bank_ad = vector.insert %carry_ad11 into %built_ad10[%c1, %c1] : i32, vector<2x2xi32> + scf.for %ks = [%c0 to %c2 step %c1] unroll { + %stage_kbase = index.mul %ks, %c32 : index + %stage_klo = index.add %stage_kbase, %weight_klo : index + %stage_khi = index.add %stage_kbase, %weight_khi : index + %stage_load_k = index.add %stage_kbase, %load_k : index + scf.for %slot = [%c0 to %c4 step %c1] unroll { + %local_row0 = index.madd %slot, %c16, %load_row : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %payload_word = vector.extract %bank_wp[%ks, %slot] : vector<2x4xi32> -> i32 + %payload_vector = vector.splat %payload_word : vector<1xi32> + %payload = vector.bitcast %payload_vector : vector<1xi32> to vector<4xi8> + %exponent_word = vector.extract %bank_we[%ks, %slot] : vector<2x4xi32> -> i32 + %exponent = scalar.trunci %exponent_word : i32 to i8 + %lo, %hi, %scale = func.call @ggml_mxfp4_decode_packet(%payload, %exponent) : (vector<4xi8>, i8) -> (vector<4xi8>, vector<4xi8>, f32) + vector.store %lo, %weight_stage_view[%local_row, %stage_klo] : vector<4xi8>, view<64x80xi8> + vector.store %hi, %weight_stage_view[%local_row, %stage_khi] : vector<4xi8>, view<64x80xi8> + scf.if %scale_writer { + %scale_tile = index.div %local_row, %c16 : index + %scale_inner = index.rem %local_row, %c16 : index + %scale_parity = index.rem %scale_inner, %c2 : index + %scale_lane = index.div %scale_inner, %c2 : index + %scale_tile_base = index.mul %scale_tile, %c16 : index + %scale_parity_base = index.madd %scale_parity, %c8, %scale_tile_base : index + %scale_slot0 = index.add %scale_parity_base, %scale_lane : index + %scale_slot = index.assume %scale_slot0 [range(%scale_slot0, 0, 63)] : index + view.store %scale, %ws[%ks, %scale_slot] : f32, view<2x64xf32> + } + } + scf.for %slot = [%c0 to %c2 step %c1] unroll { + %row0 = index.madd %slot, %c16, %load_row : index + %row = index.assume %row0 [range(%row0, 0, 31)] : index + %av_words = vector.extract %bank_ap[%ks, %slot] : vector<2x2x2xi32> -> vector<2xi32> + %av = vector.bitcast %av_words : vector<2xi32> to vector<8xi8> + %dh_word = vector.extract %bank_ad[%ks, %slot] : vector<2x2xi32> -> i32 + %dh_vector = vector.splat %dh_word : vector<1xi32> + %dh_halves = vector.bitcast %dh_vector : vector<1xi32> to vector<2xf16> + %dh = vector.extract %dh_halves[%c0] : vector<2xf16> -> f16 + %ad = scalar.extf %dh : f16 to f32 + vector.store %av, %activation_stage_physical_view[%row, %stage_load_k] : vector<8xi8>, view<32x80xi8> + scf.if %scale_writer { + view.store %ad, %xs[%ks, %row] : f32, view<2x32xf32> + } + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_block = index.add %block_base, %c2 : index + %has_next = index.cmp ult, %next_block, %quant_groups : index + %safe_next_block = scf.select %has_next, %next_block, %c0 : index + %next_wp00, %next_wp01, %next_wp02, %next_wp03, %next_wp10, %next_wp11, %next_wp12, %next_wp13, %next_we00, %next_we01, %next_we02, %next_we03, %next_we10, %next_we11, %next_we12, %next_we13, %next_ap000, %next_ap001, %next_ap010, %next_ap011, %next_ap100, %next_ap101, %next_ap110, %next_ap111, %next_ad00, %next_ad01, %next_ad10, %next_ad11 = func.call @ggml_mxfp4_prefetch_panel(%weight_noalias, %input_noalias, %route_stage, %expert_byte_base, %weight_row_bytes, %channel_tile_base, %bounded_output_size, %quant_groups, %load_row, %load_k, %weight_klo, %bounded_route_count, %bounded_input_route_count, %safe_next_block) : (buffer, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index) -> (i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32) + %next00, %next01, %next10, %next11 = scf.for %ks = [%c0 to %c2 step %c1](%part00 = %acc00 : vector<8xf32>, %part01 = %acc01 : vector<8xf32>, %part10 = %acc10 : vector<8xf32>, %part11 = %acc11 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) unroll { + %iz = vector.fragment %zero_integer shape [%m, %n] : vector<8xi32> + %i00, %i01, %i10, %i11 = scf.for %kh = [%c0 to %c32 step %c16](%v00 = %iz : vector<8xi32>, %v01 = %iz : vector<8xi32>, %v10 = %iz : vector<8xi32>, %v11 = %iz : vector<8xi32>) -> (vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>) unroll { + %stage_k = index.madd %ks, %c32, %kh : index + %a0raw = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %stage_k] shape [%m, %n] : view<64x80xi8> -> vector<16xi8> + %a0words = vector.bitcast %a0raw : vector<16xi8> to vector<4xi32> + %a0 = vector.fragment %a0words shape [%m, %n] using {schema = %i8_schema : encoding} : vector<4xi32> + %a1raw = vector.fragment.load %weight_stage_view[%subgroup_channel1, %stage_k] shape [%m, %n] : view<64x80xi8> -> vector<16xi8> + %a1words = vector.bitcast %a1raw : vector<16xi8> to vector<4xi32> + %a1 = vector.fragment %a1words shape [%m, %n] using {schema = %i8_schema : encoding} : vector<4xi32> + %b0raw = vector.fragment.load %activation_fragment_view[%stage_k, %c0] shape [%m, %n] : view<64x32xi8, %activation_fragment_layout> -> vector<16xi8> + %b0words = vector.bitcast %b0raw : vector<16xi8> to vector<4xi32> + %b0 = vector.fragment %b0words shape [%m, %n] using {schema = %i8_schema : encoding} : vector<4xi32> + %b1raw = vector.fragment.load %activation_fragment_view[%stage_k, %c16] shape [%m, %n] : view<64x32xi8, %activation_fragment_layout> -> vector<16xi8> + %b1words = vector.bitcast %b1raw : vector<16xi8> to vector<4xi32> + %b1 = vector.fragment %b1words shape [%m, %n] using {schema = %i8_schema : encoding} : vector<4xi32> + %d00 = vector.mma %a0, %b0, %v00 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %d01 = vector.mma %a0, %b1, %v01 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %d10 = vector.mma %a1, %b0, %v10 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %d11 = vector.mma %a1, %b1, %v11 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.yield %d00, %d01, %d10, %d11 : vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32> + } + %wsi00 = index.add %subgroup_channel_add, %lane_row_add : index + %wsi0 = index.assume %wsi00 [range(%wsi00, 0, 56), multiple_of(%wsi00, 8)] : index + %wd0 = vector.load %ws[%ks, %wsi0] : view<2x64xf32> -> vector<8xf32> + %wsi10 = index.add %subgroup_channel1, %lane_row_add : index + %wsi1 = index.assume %wsi10 [range(%wsi10, 0, 56), multiple_of(%wsi10, 8)] : index + %wd1 = vector.load %ws[%ks, %wsi1] : view<2x64xf32> -> vector<8xf32> + %asi00 = index.add %lane_lo, %c0 : index + %asi0 = index.assume %asi00 [range(%asi00, 0, 31)] : index + %xd0 = view.load %xs[%ks, %asi0] : view<2x32xf32> -> f32 + %xdv0 = vector.splat %xd0 : vector<8xf32> + %asi10 = index.add %lane_lo, %c16 : index + %asi1 = index.assume %asi10 [range(%asi10, 0, 31)] : index + %xd1 = view.load %xs[%ks, %asi1] : view<2x32xf32> -> f32 + %xdv1 = vector.splat %xd1 : vector<8xf32> + %prod00 = vector.mulf %wd0, %xdv0 : vector<8xf32> + %f00 = vector.sitofp %i00 : vector<8xi32> to vector<8xf32> + %n00 = vector.fmaf %prod00, %f00, %part00 : vector<8xf32> + %prod01 = vector.mulf %wd0, %xdv1 : vector<8xf32> + %f01 = vector.sitofp %i01 : vector<8xi32> to vector<8xf32> + %n01 = vector.fmaf %prod01, %f01, %part01 : vector<8xf32> + %prod10 = vector.mulf %wd1, %xdv0 : vector<8xf32> + %f10 = vector.sitofp %i10 : vector<8xi32> to vector<8xf32> + %n10 = vector.fmaf %prod10, %f10, %part10 : vector<8xf32> + %prod11 = vector.mulf %wd1, %xdv1 : vector<8xf32> + %f11 = vector.sitofp %i11 : vector<8xi32> to vector<8xf32> + %n11 = vector.fmaf %prod11, %f11, %part11 : vector<8xf32> + scf.yield %n00, %n01, %n10, %n11 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11, %next_wp00, %next_wp01, %next_wp02, %next_wp03, %next_wp10, %next_wp11, %next_wp12, %next_wp13, %next_we00, %next_we01, %next_we02, %next_we03, %next_we10, %next_we11, %next_we12, %next_we13, %next_ap000, %next_ap001, %next_ap010, %next_ap011, %next_ap100, %next_ap101, %next_ap110, %next_ap111, %next_ad00, %next_ad01, %next_ad10, %next_ad11 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32 + } + %publish_route0 = index.div %lane, %c2 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 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 + %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<8xf32>, view<16x16xf32, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %wide = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + template.apply<@ggml.matmul_id.mxfp4_prefetch.publish_vector8>(%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<8xf32>, 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<8xf32>, view<16x16xf32, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %wide = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + template.apply<@ggml.matmul_id.mxfp4_prefetch.publish_vector8>(%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<8xf32>, 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<8xf32>, view<16x16xf32, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %wide = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + template.apply<@ggml.matmul_id.mxfp4_prefetch.publish_vector8>(%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<8xf32>, 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<8xf32>, view<16x16xf32, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %wide = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + template.apply<@ggml.matmul_id.mxfp4_prefetch.publish_vector8>(%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<8xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul_id.mxfp4_prefetch.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_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom index 4a7648cd0dd0..99112a9fd18f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom @@ -56,11 +56,11 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %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 + %wave_result_stage_bytes = index.constant 1024 : offset + %result_stage_bytes = index.constant 2048 : 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> + %zero_accumulator = vector.constant 0.0 : vector<4xf32> %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index @@ -87,8 +87,8 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %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> + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf32> %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 @@ -97,14 +97,14 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %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> + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<4xf32> %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>) { + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<4xf32>, %block_acc01 = %init01 : vector<4xf32>, %block_acc10 = %init10 : vector<4xf32>, %block_acc11 = %init11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<4xf32>, %acc01 = %block_acc01 : vector<4xf32>, %acc10 = %block_acc10 : vector<4xf32>, %acc11 = %block_acc11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { %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 @@ -145,21 +145,21 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c } } 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 { + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<4xf32>, %half_acc01 = %acc01 : vector<4xf32>, %half_acc10 = %acc10 : vector<4xf32>, %half_acc11 = %acc11 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) 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> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<4xf32> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %next00, %next01, %next10, %next11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> } %publish_token0 = index.div %lane, %c4 : index %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index @@ -181,35 +181,31 @@ func.def inline @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_c %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> + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, 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> + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, 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> + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, 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> + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<4xf32>, view<16x16xf32, %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> + %wide = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<4xf32> template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) } template.apply<@ggml.mul_mat_f32_f32_wmma.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) 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 index 59eaa9d8c006..30d991b9eff5 100644 --- 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 @@ -7,6 +7,29 @@ template.decl @ggml.rows.cache_publish.f32_vector2(%output_format: index, %publi 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.decl @ggml.rows.cache_publish.safe_i32_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) + +template.def<@ggml.rows.cache_publish.safe_i32_row> device @ggml_rows_cache_publish_safe_i32_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_i32 = scalar.constant 0 : i32 + %c1048575_i32 = scalar.constant 1048575 : i32 + %c0_offset = index.constant 0 : offset + %indices_view = buffer.view %indices[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %index_raw = view.load %indices_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i32 : i32 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i32 : i32 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i32 = scf.select %valid_index, %index_raw, %c0_i32 : i32 + %safe_index_i32 = scalar.assume %safe_index0_i32 [range(%safe_index0_i32, 0, 1048575)] : i32 + %cache_row0 = index.cast %safe_index_i32 : i32 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.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 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_packed.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_packed.loom new file mode 100644 index 000000000000..24aaa97d8ee9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_packed.loom @@ -0,0 +1,76 @@ +config.decl @ggml.concat_packed.lhs_chunk : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.concat_packed.rhs_chunk : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.concat_packed.lhs_span : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.concat_packed.rhs_span : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.concat_packed.element_bits : %v: index where [range(%v, 16, 32)] + +kernel.def @ggml_concat_packed(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 134217727), lt(%linear0, %count)] : index, index + %zero = index.constant 0 : offset + %lc = config.get @ggml.concat_packed.lhs_chunk : index + %rc = config.get @ggml.concat_packed.rhs_chunk : index + %ls = config.get @ggml.concat_packed.lhs_span : index + %rs = config.get @ggml.concat_packed.rhs_span : index + %chunk = index.add %lc, %rc : index + %within = index.rem %linear, %chunk : index + %batch = index.div %linear, %chunk : index + %is_left = index.cmp ult, %within, %lc : index + %bits = config.get @ggml.concat_packed.element_bits : index + %bits16 = index.constant 16 : index + %is16 = index.cmp eq, %bits, %bits16 : index + scf.if %is16 { + %yv = buffer.view %output[%zero] : buffer -> view<[%count]xi16> + %value = scf.if %is_left -> (i16) { + %pos0 = index.madd %batch, %lc, %within : index + %pos, %span = index.assume %pos0, %ls [range(%pos0, 0, 134217727), lt(%pos0, %ls)] : index, index + %lv = buffer.view %lhs[%zero] : buffer -> view<[%ls]xi16> + %v = view.load %lv[%pos] : view<[%ls]xi16> -> i16 + scf.yield %v : i16 + } else { + %right = index.sub %within, %lc : index + %pos0 = index.madd %batch, %rc, %right : index + %pos, %span = index.assume %pos0, %rs [range(%pos0, 0, 134217727), lt(%pos0, %rs)] : index, index + %rv = buffer.view %rhs[%zero] : buffer -> view<[%rs]xi16> + %v = view.load %rv[%pos] : view<[%rs]xi16> -> i16 + scf.yield %v : i16 + } + view.store %value, %yv[%linear] : i16, view<[%count]xi16> + } else { + %yv = buffer.view %output[%zero] : buffer -> view<[%count]xi32> + %value = scf.if %is_left -> (i32) { + %pos0 = index.madd %batch, %lc, %within : index + %pos, %span = index.assume %pos0, %ls [range(%pos0, 0, 134217727), lt(%pos0, %ls)] : index, index + %lv = buffer.view %lhs[%zero] : buffer -> view<[%ls]xi32> + %v = view.load %lv[%pos] : view<[%ls]xi32> -> i32 + scf.yield %v : i32 + } else { + %right = index.sub %within, %lc : index + %pos0 = index.madd %batch, %rc, %right : index + %pos, %span = index.assume %pos0, %rs [range(%pos0, 0, 134217727), lt(%pos0, %rs)] : index, index + %rv = buffer.view %rhs[%zero] : buffer -> view<[%rs]xi32> + %v = view.load %rv[%pos] : view<[%rs]xi32> -> i32 + scf.yield %v : i32 + } + view.store %value, %yv[%linear] : i32, view<[%count]xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_bits.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_bits.loom new file mode 100644 index 000000000000..9693d401faf9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_bits.loom @@ -0,0 +1,73 @@ +config.decl @ggml.copy_bits.source_ne0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.source_ne1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.source_ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.source_ne3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.ne0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.ne1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.stride3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_bits.element_bits : %value: index where [range(%value, 1, 1073741824)] +kernel.def @ggml_copy_bits(%element_count: index, %source_span: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %source_span: index, %source: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 1073741824)] : index + %span = index.assume %source_span [range(%source_span, 1, 1073741824)] : index + %ne0 = config.get @ggml.copy_bits.ne0 : index + %ne1 = config.get @ggml.copy_bits.ne1 : index + %ne2 = config.get @ggml.copy_bits.ne2 : index + %stride0 = config.get @ggml.copy_bits.stride0 : index + %stride1 = config.get @ggml.copy_bits.stride1 : index + %stride2 = config.get @ggml.copy_bits.stride2 : index + %stride3 = config.get @ggml.copy_bits.stride3 : index + %element_bits = config.get @ggml.copy_bits.element_bits : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 1073741823), lt(%linear0, %count)] : index, index + %i0 = index.rem %linear, %ne0 : index + %rest0 = index.div %linear, %ne0 : index + %i1 = index.rem %rest0, %ne1 : index + %rest1 = index.div %rest0, %ne1 : index + %i2 = index.rem %rest1, %ne2 : index + %i3 = index.div %rest1, %ne2 : index + %sn0 = config.get @ggml.copy_bits.source_ne0 : index + %si0 = index.rem %i0, %sn0 : index + %sn1 = config.get @ggml.copy_bits.source_ne1 : index + %si1 = index.rem %i1, %sn1 : index + %sn2 = config.get @ggml.copy_bits.source_ne2 : index + %si2 = index.rem %i2, %sn2 : index + %sn3 = config.get @ggml.copy_bits.source_ne3 : index + %si3 = index.rem %i3, %sn3 : index + %p0 = index.mul %si0, %stride0 : index + %p1 = index.madd %si1, %stride1, %p0 : index + %p2 = index.madd %si2, %stride2, %p1 : index + %p3 = index.madd %si3, %stride3, %p2 : index + %p, %source_extent = index.assume %p3, %span [range(%p3, 0, 1073741823), lt(%p3, %span)] : index, index + %zero = index.constant 0 : offset + %bits16 = index.constant 16 : index + %is16 = index.cmp eq, %element_bits, %bits16 : index + scf.if %is16 { + %src = buffer.view %source[%zero] : buffer -> view<[%span]xi16> + %dst = buffer.view %output[%zero] : buffer -> view<[%count]xi16> + %value = view.load %src[%p] : view<[%span]xi16> -> i16 + view.store %value, %dst[%linear] : i16, view<[%count]xi16> + } else { + %src = buffer.view %source[%zero] : buffer -> view<[%span]xi32> + %dst = buffer.view %output[%zero] : buffer -> view<[%count]xi32> + %value = view.load %src[%p] : view<[%span]xi32> -> i32 + view.store %value, %dst[%linear] : i32, view<[%count]xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/add_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/add_f16.loom new file mode 100644 index 000000000000..55baf9c80be1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/add_f16.loom @@ -0,0 +1,30 @@ +kernel.def @ggml_add_f16(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 134217727), lt(%linear0, %count)] : index, index + %zero = index.constant 0 : offset + %lv = buffer.view %lhs[%zero] : buffer -> view<[%count]xf16> + %rv = buffer.view %rhs[%zero] : buffer -> view<[%count]xf16> + %yv = buffer.view %output[%zero] : buffer -> view<[%count]xf16> + %l = view.load %lv[%linear] : view<[%count]xf16> -> f16 + %r = view.load %rv[%linear] : view<[%count]xf16> -> f16 + %lf = scalar.extf %l : f16 to f32 + %rf = scalar.extf %r : f16 to f32 + %sum = scalar.addf %lf, %rf : f32 + %result = scalar.fptrunc %sum : f32 to f16 + view.store %result, %yv[%linear] : f16, view<[%count]xf16> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom index f8e74b98b4bd..fb64b6ba53ce 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom @@ -13,17 +13,21 @@ config.decl @ggml.binary_f32.ne1 : %value: index where [range(%value, 1, 1342177 config.decl @ggml.binary_f32.ne2 : %value: index where [range(%value, 1, 134217728)] -config.decl @ggml.binary_f32.src0_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src0_stride0 : %value: index where [range(%value, 0, 1)] -config.decl @ggml.binary_f32.src0_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src0_stride1 : %value: index where [range(%value, 0, 134217728)] -config.decl @ggml.binary_f32.src0_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src0_stride2 : %value: index where [range(%value, 0, 134217728)] -config.decl @ggml.binary_f32.src1_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src0_stride3 : %value: index where [range(%value, 0, 134217728)] -config.decl @ggml.binary_f32.src1_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src1_stride0 : %value: index where [range(%value, 0, 1)] -config.decl @ggml.binary_f32.src1_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.binary_f32.src1_stride1 : %value: index where [range(%value, 0, 134217728)] + +config.decl @ggml.binary_f32.src1_stride2 : %value: index where [range(%value, 0, 134217728)] + +config.decl @ggml.binary_f32.src1_stride3 : %value: index where [range(%value, 0, 134217728)] config.decl @ggml.binary_f32.src0_span : %value: index where [range(%value, 1, 134217728)] @@ -41,9 +45,11 @@ kernel.def target(@ggml_binary_f32_gfx11_wave64) export("ggml_binary_f32") @ggml %n0 = config.get @ggml.binary_f32.ne0 : index %n1 = config.get @ggml.binary_f32.ne1 : index %n2 = config.get @ggml.binary_f32.ne2 : index + %s0_0 = config.get @ggml.binary_f32.src0_stride0 : index %s0_1 = config.get @ggml.binary_f32.src0_stride1 : index %s0_2 = config.get @ggml.binary_f32.src0_stride2 : index %s0_3 = config.get @ggml.binary_f32.src0_stride3 : index + %s1_0 = config.get @ggml.binary_f32.src1_stride0 : index %s1_1 = config.get @ggml.binary_f32.src1_stride1 : index %s1_2 = config.get @ggml.binary_f32.src1_stride2 : index %s1_3 = config.get @ggml.binary_f32.src1_stride3 : index @@ -69,16 +75,18 @@ kernel.def target(@ggml_binary_f32_gfx11_wave64) export("ggml_binary_f32") @ggml %linear_div_n01 = index.div %linear_div_n0, %n1 : index %coord2 = index.rem %linear_div_n01, %n2 : index %coord3 = index.div %linear_div_n01, %n2 : index + %src0_term0 = index.mul %coord0, %s0_0 : index %src0_term1 = index.mul %coord1, %s0_1 : index %src0_term2 = index.mul %coord2, %s0_2 : index %src0_term3 = index.mul %coord3, %s0_3 : index - %src0_linear01 = index.add %coord0, %src0_term1 : index + %src0_linear01 = index.add %src0_term0, %src0_term1 : index %src0_linear012 = index.add %src0_linear01, %src0_term2 : index %src0_linear = index.add %src0_linear012, %src0_term3 : index + %src1_term0 = index.mul %coord0, %s1_0 : index %src1_term1 = index.mul %coord1, %s1_1 : index %src1_term2 = index.mul %coord2, %s1_2 : index %src1_term3 = index.mul %coord3, %s1_3 : index - %src1_linear01 = index.add %coord0, %src1_term1 : index + %src1_linear01 = index.add %src1_term0, %src1_term1 : index %src1_linear012 = index.add %src1_linear01, %src1_term2 : index %src1_linear = index.add %src1_linear012, %src1_term3 : index %src0_in_bounds = index.cmp ult, %src0_linear, %src0_count : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/clamp_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/clamp_f32.loom new file mode 100644 index 000000000000..66bfdce17b85 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/clamp_f32.loom @@ -0,0 +1,35 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_clamp_f32_target {subgroup_size = 64} +config.decl @ggml.clamp.min : f32 +config.decl @ggml.clamp.max : f32 + +kernel.def target(@ggml_clamp_f32_target) export("ggml_clamp_f32") @ggml_clamp_f32(%element_count: index) { + %one = index.constant 1 : index + %block = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %groups = index.div %rounded, %block : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%block, %one, %one) : index +} launch(%element_count: index, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %block = index.constant 256 : index + %base = index.mul %group, %block : index + %linear = index.add %base, %thread : index + %valid = index.cmp ult, %linear, %count : index + %lower = config.get @ggml.clamp.min : f32 + %upper = config.get @ggml.clamp.max : f32 + %zero = index.constant 0 : offset + %src = buffer.view %input[%zero] : buffer -> view<[%count]xf32> + %dst = buffer.view %output[%zero] : buffer -> view<[%count]xf32> + scf.if %valid { + %value = view.load %src[%linear] : view<[%count]xf32> -> f32 + %bounded_lower = scalar.maxnumf %value, %lower : f32 + %bounded = scalar.minnumf %bounded_lower, %upper : f32 + view.store %bounded, %dst[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/fill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/fill.loom new file mode 100644 index 000000000000..63f817328942 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/fill.loom @@ -0,0 +1,34 @@ +config.decl @ggml.fill.value : f32 +config.decl @ggml.fill.format : %value: index where [range(%value, 16, 32)] +kernel.def @ggml_fill(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 134217727), lt(%linear0, %count)] : index, index + %zero = index.constant 0 : offset + %format = config.get @ggml.fill.format : index + %bits16 = index.constant 16 : index + %is16 = index.cmp eq, %format, %bits16 : index + %value = config.get @ggml.fill.value : f32 + scf.if %is16 { + %v = scalar.fptrunc %value : f32 to f16 + %out = buffer.view %output[%zero] : buffer -> view<[%count]xf16> + view.store %v, %out[%linear] : f16, view<[%count]xf16> + } else { + %out = buffer.view %output[%zero] : buffer -> view<[%count]xf32> + view.store %value, %out[%linear] : f32, view<[%count]xf32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_comb_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_comb_f32.loom new file mode 100644 index 000000000000..17896c33fb69 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_comb_f32.loom @@ -0,0 +1,120 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_hc_comb_target {subgroup_size = 64} +config.decl @ggml.hc.tokens : %value: index where [range(%value, 1, 32768)] +config.decl @ggml.hc.n_iter : %value: index where [range(%value, 1, 128)] +config.decl @ggml.hc.eps : f32 + +func.def inline @hc_normalize_row(%row: vector<4xf32>, %eps: f32) -> (vector<4xf32>) { + %sum = vector.reduce %row, %eps : vector<4xf32>, f32 + %one = scalar.constant 1.0 : f32 + %inverse = scalar.divf %one, %sum : f32 + %inverse_vector = vector.splat %inverse : vector<4xf32> + %normalized = vector.mulf %row, %inverse_vector : vector<4xf32> + func.return %normalized : vector<4xf32> +} + +func.def inline @hc_normalize_cols(%r0: vector<4xf32>, %r1: vector<4xf32>, %r2: vector<4xf32>, %r3: vector<4xf32>, %eps: f32) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %eps_vector = vector.splat %eps : vector<4xf32> + %s0 = vector.addf %eps_vector, %r0 : vector<4xf32> + %s1 = vector.addf %s0, %r1 : vector<4xf32> + %s2 = vector.addf %s1, %r2 : vector<4xf32> + %sum = vector.addf %s2, %r3 : vector<4xf32> + %ones = vector.constant 1.0 : vector<4xf32> + %inverse = vector.divf %ones, %sum : vector<4xf32> + %n0 = vector.mulf %r0, %inverse : vector<4xf32> + %n1 = vector.mulf %r1, %inverse : vector<4xf32> + %n2 = vector.mulf %r2, %inverse : vector<4xf32> + %n3 = vector.mulf %r3, %inverse : vector<4xf32> + func.return %n0, %n1, %n2, %n3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +func.def inline @hc_softmax_row(%mix: vector<4xf32>, %scale: vector<4xf32>, %base: vector<4xf32>, %eps: f32) -> (vector<4xf32>) { + %scaled = vector.mulf %mix, %scale : vector<4xf32> + %affine = vector.addf %scaled, %base : vector<4xf32> + %negative_inf = scalar.constant -inf : f32 + %maximum = vector.reduce %affine, %negative_inf : vector<4xf32>, f32 + %maximum_vector = vector.splat %maximum : vector<4xf32> + %shifted = vector.subf %affine, %maximum_vector : vector<4xf32> + %exponential = vector.expf %shifted : vector<4xf32> + %zero = scalar.constant 0.0 : f32 + %probability = func.call @hc_normalize_row(%exponential, %zero) : (vector<4xf32>, f32) -> (vector<4xf32>) + %eps_vector = vector.splat %eps : vector<4xf32> + %result = vector.addf %probability, %eps_vector : vector<4xf32> + func.return %result : vector<4xf32> +} + +kernel.def target(@ggml_hc_comb_target) export("ggml_hc_comb_f32") @ggml_hc_comb_f32() { + %one = index.constant 1 : index + %block = index.constant 64 : index + %rounding = index.constant 63 : index + %tokens = config.get @ggml.hc.tokens : index + %rounded = index.add %tokens, %rounding : index + %groups = index.div %rounded, %block : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%block, %one, %one) : index +} launch(%mixes: buffer, %scale: buffer, %base: buffer, %output: buffer) { + %tokens = config.get @ggml.hc.tokens : index + %iterations = config.get @ggml.hc.n_iter : index + %eps = config.get @ggml.hc.eps : f32 + %one = index.constant 1 : index + %two = index.constant 2 : index + %four = index.constant 4 : index + %eight = index.constant 8 : index + %twelve = index.constant 12 : index + %sixteen = index.constant 16 : index + %twenty = index.constant 20 : index + %twentyfour = index.constant 24 : index + %zero_offset = index.constant 0 : offset + %mix_count = index.mul %tokens, %twentyfour : index + %out_count = index.mul %tokens, %sixteen : index + %m = buffer.view %mixes[%zero_offset] : buffer -> view<[%mix_count]xf32> + %s = buffer.view %scale[%zero_offset] : buffer -> view<3xf32> + %b = buffer.view %base[%zero_offset] : buffer -> view<24xf32> + %y = buffer.view %output[%zero_offset] : buffer -> view<[%out_count]xf32> + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %block = index.constant 64 : index + %group_base = index.mul %group, %block : index + %token = index.add %group_base, %thread : index + %valid = index.cmp ult, %token, %tokens : index + scf.if %valid { + %scale_value = view.load %s[%two] : view<3xf32> -> f32 + %scale_vector = vector.splat %scale_value : vector<4xf32> + %mix_base = index.mul %token, %twentyfour : index + %m0i = index.add %mix_base, %eight : index + %m1i = index.add %mix_base, %twelve : index + %m2i = index.add %mix_base, %sixteen : index + %m3i = index.add %mix_base, %twenty : index + %m0 = vector.load %m[%m0i] : view<[%mix_count]xf32> -> vector<4xf32> + %m1 = vector.load %m[%m1i] : view<[%mix_count]xf32> -> vector<4xf32> + %m2 = vector.load %m[%m2i] : view<[%mix_count]xf32> -> vector<4xf32> + %m3 = vector.load %m[%m3i] : view<[%mix_count]xf32> -> vector<4xf32> + %b0 = vector.load %b[%eight] : view<24xf32> -> vector<4xf32> + %b1 = vector.load %b[%twelve] : view<24xf32> -> vector<4xf32> + %b2 = vector.load %b[%sixteen] : view<24xf32> -> vector<4xf32> + %b3 = vector.load %b[%twenty] : view<24xf32> -> vector<4xf32> + %r0 = func.call @hc_softmax_row(%m0, %scale_vector, %b0, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>) + %r1 = func.call @hc_softmax_row(%m1, %scale_vector, %b1, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>) + %r2 = func.call @hc_softmax_row(%m2, %scale_vector, %b2, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>) + %r3 = func.call @hc_softmax_row(%m3, %scale_vector, %b3, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>) + %c0, %c1, %c2, %c3 = func.call @hc_normalize_cols(%r0, %r1, %r2, %r3, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %result0, %result1, %result2, %result3 = scf.for %iter = [%one to %iterations step %one](%a0 = %c0 : vector<4xf32>, %a1 = %c1 : vector<4xf32>, %a2 = %c2 : vector<4xf32>, %a3 = %c3 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %n0 = func.call @hc_normalize_row(%a0, %eps) : (vector<4xf32>, f32) -> (vector<4xf32>) + %n1 = func.call @hc_normalize_row(%a1, %eps) : (vector<4xf32>, f32) -> (vector<4xf32>) + %n2 = func.call @hc_normalize_row(%a2, %eps) : (vector<4xf32>, f32) -> (vector<4xf32>) + %n3 = func.call @hc_normalize_row(%a3, %eps) : (vector<4xf32>, f32) -> (vector<4xf32>) + %next0, %next1, %next2, %next3 = func.call @hc_normalize_cols(%n0, %n1, %n2, %n3, %eps) : (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, f32) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + scf.yield %next0, %next1, %next2, %next3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %y0i = index.mul %token, %sixteen : index + %y1i = index.add %y0i, %four : index + %y2i = index.add %y0i, %eight : index + %y3i = index.add %y0i, %twelve : index + vector.store %result0, %y[%y0i] : vector<4xf32>, view<[%out_count]xf32> + vector.store %result1, %y[%y1i] : vector<4xf32>, view<[%out_count]xf32> + vector.store %result2, %y[%y2i] : vector<4xf32>, view<[%out_count]xf32> + vector.store %result3, %y[%y3i] : vector<4xf32>, view<[%out_count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_post_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_post_f32.loom new file mode 100644 index 000000000000..365b6c49b349 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_post_f32.loom @@ -0,0 +1,87 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_hc_post_target {subgroup_size = 64} + +config.decl @ggml.hc.tokens : %value: index where [range(%value, 1, 32768)] + +config.decl @ggml.hc.embd : %value: index where [range(%value, 1, 32768)] + +config.decl @ggml.hc.streams : %value: index where [range(%value, 1, 128)] + +config.decl @ggml.hc.identity : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_hc_post_target) export("ggml_hc_post_f32") @ggml_hc_post_f32() { + %one = index.constant 1 : index + %streams = config.get @ggml.hc.streams : index + %block = index.constant 256 : index + %rounding = index.constant 255 : index + %tokens = config.get @ggml.hc.tokens : index + %embd = config.get @ggml.hc.embd : index + %rows = index.mul %tokens, %streams : index + %count = index.mul %rows, %embd : index + %rounded = index.add %count, %rounding : index + %groups = index.div %rounded, %block : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%block, %one, %one) : index +} launch(%input: buffer, %residual: buffer, %post: buffer, %comb: buffer, %output: buffer) { + %tokens = config.get @ggml.hc.tokens : index + %embd = config.get @ggml.hc.embd : index + %one = index.constant 1 : index + %streams = config.get @ggml.hc.streams : index + %zero = index.constant 0 : index + %zero_offset = index.constant 0 : offset + %input_count = index.mul %tokens, %embd : index + %count = index.mul %input_count, %streams : index + %post_count = index.mul %tokens, %streams : index + %x = buffer.view %input[%zero_offset] : buffer -> view<[%input_count]xf32> + %r = buffer.view %residual[%zero_offset] : buffer -> view<[%count]xf32> + %p = buffer.view %post[%zero_offset] : buffer -> view<[%post_count]xf32> + %y = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %block = index.constant 256 : index + %base = index.mul %group, %block : index + %linear = index.add %base, %thread : index + %valid = index.cmp ult, %linear, %count : index + scf.if %valid { + %feature = index.rem %linear, %embd : index + %row = index.div %linear, %embd : index + %dst = index.rem %row, %streams : index + %token = index.div %row, %streams : index + %x_base = index.mul %token, %embd : index + %x_index = index.add %x_base, %feature : index + %p_base = index.mul %token, %streams : index + %p_index = index.add %p_base, %dst : index + %x_value = view.load %x[%x_index] : view<[%input_count]xf32> -> f32 + %p_value = view.load %p[%p_index] : view<[%post_count]xf32> -> f32 + %initial = scalar.mulf %x_value, %p_value : f32 + %identity_mode = config.get @ggml.hc.identity : index + %identity = index.cmp ne, %identity_mode, %zero : index + %sum = scf.if %identity -> (f32) { + %residual_value = view.load %r[%linear] : view<[%count]xf32> -> f32 + %result = scalar.addf %initial, %residual_value : f32 + scf.yield %result : f32 + } else { + %comb_count = index.mul %post_count, %streams : index + %c = buffer.view %comb[%zero_offset] : buffer -> view<[%comb_count]xf32> + %r_base = index.mul %p_base, %embd : index + %c_base = index.mul %p_base, %streams : index + %c_dst = index.add %c_base, %dst : index + %mixed = scf.for %h = [%zero to %streams step %one](%acc = %initial : f32) -> (f32) unroll { + %r_offset = index.mul %h, %embd : index + %r_row = index.add %r_base, %r_offset : index + %r_index = index.add %r_row, %feature : index + %c_offset = index.mul %h, %streams : index + %c_index = index.add %c_dst, %c_offset : index + %r_value = view.load %r[%r_index] : view<[%count]xf32> -> f32 + %c_value = view.load %c[%c_index] : view<[%comb_count]xf32> -> f32 + %term = scalar.mulf %r_value, %c_value : f32 + %next = scalar.addf %acc, %term : f32 + scf.yield %next : f32 + } + scf.yield %mixed : f32 + } + view.store %sum, %y[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_pre_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_pre_f32.loom new file mode 100644 index 000000000000..979935ef0c5f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/hc_pre_f32.loom @@ -0,0 +1,89 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_hc_pre_target {subgroup_size = 64} + +config.decl @ggml.hc.tokens : %value: index where [range(%value, 1, 32768)] + +config.decl @ggml.hc.embd : %value: index where [range(%value, 1, 32768)] + +config.decl @ggml.hc.streams : %value: index where [range(%value, 1, 128)] + +config.decl @ggml.hc.gated : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.hc.scale : f32 + +kernel.def target(@ggml_hc_pre_target) export("ggml_hc_pre_f32") @ggml_hc_pre_f32() { + %one = index.constant 1 : index + %block = index.constant 256 : index + %rounding = index.constant 255 : index + %tokens = config.get @ggml.hc.tokens : index + %embd = config.get @ggml.hc.embd : index + %count = index.mul %tokens, %embd : index + %rounded = index.add %count, %rounding : index + %groups = index.div %rounded, %block : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%block, %one, %one) : index +} launch(%input: buffer, %weights: buffer, %output: buffer) { + %tokens = config.get @ggml.hc.tokens : index + %embd = config.get @ggml.hc.embd : index + %one = index.constant 1 : index + %streams = config.get @ggml.hc.streams : index + %zero = index.constant 0 : index + %zero_offset = index.constant 0 : offset + %zero_f = scalar.constant 0.0 : f32 + %count = index.mul %tokens, %embd : index + %input_count = index.mul %count, %streams : index + %gated_mode = config.get @ggml.hc.gated : index + %gated = index.cmp ne, %gated_mode, %zero : index + %weight_count = scf.if %gated -> (index) { + scf.yield %input_count : index + } else { + %weights = index.mul %tokens, %streams : index + scf.yield %weights : index + } + %x = buffer.view %input[%zero_offset] : buffer -> view<[%input_count]xf32> + %w = buffer.view %weights[%zero_offset] : buffer -> view<[%weight_count]xf32> + %y = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %block = index.constant 256 : index + %base = index.mul %group, %block : index + %linear = index.add %base, %thread : index + %valid = index.cmp ult, %linear, %count : index + scf.if %valid { + %feature = index.rem %linear, %embd : index + %token = index.div %linear, %embd : index + %weight_base = index.mul %token, %streams : index + %input_base = index.mul %weight_base, %embd : index + %sum = scf.for %h = [%zero to %streams step %one](%acc = %zero_f : f32) -> (f32) unroll { + %h_offset = index.mul %h, %embd : index + %row = index.add %input_base, %h_offset : index + %x_index = index.add %row, %feature : index + %w_index = scf.if %gated -> (index) { + scf.yield %x_index : index + } else { + %index = index.add %weight_base, %h : index + scf.yield %index : index + } + %x_value = view.load %x[%x_index] : view<[%input_count]xf32> -> f32 + %weight = view.load %w[%w_index] : view<[%weight_count]xf32> -> f32 + %w_value = scf.if %gated -> (f32) { + %one_f = scalar.constant 1.0 : f32 + %negative = scalar.negf %weight : f32 + %exp = scalar.expf %negative : f32 + %denominator = scalar.addf %one_f, %exp : f32 + %sigmoid = scalar.divf %one_f, %denominator : f32 + scf.yield %sigmoid : f32 + } else { + scf.yield %weight : f32 + } + %term = scalar.mulf %x_value, %w_value : f32 + %next = scalar.addf %acc, %term : f32 + scf.yield %next : f32 + } + %scale = config.get @ggml.hc.scale : f32 + %result = scalar.mulf %sum, %scale : f32 + view.store %result, %y[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/swiglu_clamp_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/swiglu_clamp_f32.loom new file mode 100644 index 000000000000..c2d58da1cb9e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/swiglu_clamp_f32.loom @@ -0,0 +1,42 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +config.decl @ggml.swiglu_clamp.limit : f32 + +kernel.def @ggml_swiglu_clamp_f32(%element_count: index) { + %one = index.constant 1 : index + %block = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %block : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%block, %one, %one) : index +} launch(%element_count: index, %gate: buffer, %up: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %block = index.constant 256 : index + %linear0 = index.madd %group, %block, %thread : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 134217727), lt(%linear0, %count)] : index, index + %zero = index.constant 0 : offset + %gate_view = buffer.view %gate[%zero] : buffer -> view<[%count]xf32> + %up_view = buffer.view %up[%zero] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output[%zero] : buffer -> view<[%count]xf32> + %limit = config.get @ggml.swiglu_clamp.limit : f32 + %lower = scalar.negf %limit : f32 + %gate_value = view.load %gate_view[%linear] : view<[%count]xf32> -> f32 + %up_value = view.load %up_view[%linear] : view<[%count]xf32> -> f32 + %bounded_gate = scalar.minnumf %gate_value, %limit : f32 + %bounded_up_lower = scalar.maxnumf %up_value, %lower : f32 + %bounded_up = scalar.minnumf %bounded_up_lower, %limit : f32 + %negative_gate = scalar.negf %bounded_gate : f32 + %exp = scalar.expf %negative_gate : f32 + %one = scalar.constant 1.0 : f32 + %denominator = scalar.addf %one, %exp : f32 + %activated = scalar.divf %bounded_gate, %denominator : f32 + %result = scalar.mulf %activated, %bounded_up : f32 + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom index b858a1650c06..808a7a6665dd 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom @@ -35,3 +35,44 @@ kernel.def target(@ggml_unary_f32_gfx11_wave64) export("ggml_unary_f32") @ggml_u } kernel.return } + +kernel.def target(@ggml_unary_f32_gfx11_wave64) export("ggml_softplus_f32") @ggml_softplus_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %input_view = buffer.view %input[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %value = view.load %input_view[%linear] : view<[%count]xf32> -> f32 + %zero = scalar.constant 0.0 : f32 + %positive = scalar.maxnumf %value, %zero : f32 + %absolute = scalar.absf %value : f32 + %negative = scalar.negf %absolute : f32 + %exp = scalar.expf %negative : f32 + %one_f32 = scalar.constant 1.0 : f32 + %sum = scalar.addf %one_f32, %exp : f32 + %delta = scalar.subf %sum, %one_f32 : f32 + %tiny = scalar.cmpf oeq, %sum, %one_f32 : f32 + %denominator = scf.select %tiny, %one_f32, %delta : f32 + %ratio = scalar.divf %exp, %denominator : f32 + %log = scalar.logf %sum : f32 + %compensated = scalar.mulf %log, %ratio : f32 + %correction = scf.select %tiny, %exp, %compensated : f32 + %result = scalar.addf %positive, %correction : f32 + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/compact_mask_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/compact_mask_f16.loom new file mode 100644 index 000000000000..a176400075cd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/compact_mask_f16.loom @@ -0,0 +1,106 @@ +amdgpu.target @ggml_compact_mask_wave64 {subgroup_size = 64} + +config.decl @ggml.sparse_attention.capacity : %value: index where [range(%value, 64, 2048), multiple_of(%value, 64)] + +// The caller bounds finite entries per row by capacity. Keep their source order. +kernel.def target(@ggml_compact_mask_wave64) @ggml_compact_mask_f16(%query_count: index, %key_count: index) { + %one = index.constant 1 : index + %threads = index.constant 256 : index + kernel.launch.config workgroups(%one, %query_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%query_count: index, %key_count: index, %mask: buffer, %indices: buffer, %compact_mask: buffer) { + %nq = index.assume %query_count [range(%query_count, 1, 2048)] : index + %nk = index.assume %key_count [range(%key_count, 1, 131072)] : index + %capacity = config.get @ggml.sparse_attention.capacity : index + %row0 = kernel.workgroup.id : index + %row = index.assume %row0 [lt(%row0, %nq)] : index + %thread = kernel.workitem.id : index + %zero = index.constant 0 : index + %step = index.constant 256 : index + %offset = index.constant 0 : offset + %zero_i32 = scalar.constant 0 : i32 + %one_i32 = scalar.constant 1 : i32 + %negative_infinity_bits = scalar.constant -1024 : i16 + %positive_infinity_bits = scalar.constant 31744 : i16 + %negative_infinity = scalar.bitcast %negative_infinity_bits : i16 to f16 + %positive_infinity = scalar.bitcast %positive_infinity_bits : i16 to f16 + %source = buffer.view %mask[%offset] : buffer -> view<[%nq]x[%nk]xf16> + %selection = buffer.view %indices[%offset] : buffer -> view<[%nq]x[%capacity]xi32> + %target = buffer.view %compact_mask[%offset] : buffer -> view<[%nq]x[%capacity]xf16> + %total = scf.for %base = [%zero to %nk step %step](%count = %zero_i32 : i32) -> (i32) { + %position = index.add %base, %thread : index + %valid = index.cmp ult, %position, %nk : index + %value = scf.if %valid -> (f16) { + %column = index.assume %position [lt(%position, %nk)] : index + %loaded = view.load %source[%row, %column] : view<[%nq]x[%nk]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %negative_infinity : f16 + } + %above = scalar.cmpf ogt, %value, %negative_infinity : f16 + %below = scalar.cmpf olt, %value, %positive_infinity : f16 + %finite = scalar.andi %above, %below : i1 + %increment = scf.select %finite, %one_i32, %zero_i32 : i32 + %prefix = kernel.workgroup.scan %increment {direction = forward, mode = exclusive} : i32 + %block_count = kernel.workgroup.reduce %increment : i32 + %slot_i32 = scalar.addi %count, %prefix : i32 + %slot0 = index.cast %slot_i32 : i32 to index + %fits = index.cmp ult, %slot0, %capacity : index + %write = scalar.andi %finite, %fits : i1 + scf.if %write { + %slot = index.assume %slot0 [lt(%slot0, %capacity)] : index + %id = index.cast %position : index to i32 + view.store %id, %selection[%row, %slot] : i32, view<[%nq]x[%capacity]xi32> + view.store %value, %target[%row, %slot] : f16, view<[%nq]x[%capacity]xf16> + } + %next_count = scalar.addi %count, %block_count : i32 + scf.yield %next_count : i32 + } + %used = index.cast %total : i32 to index + scf.for %base = [%zero to %capacity step %step] { + %slot0 = index.add %base, %thread : index + %valid = index.cmp ult, %slot0, %capacity : index + %padding = index.cmp uge, %slot0, %used : index + %write = scalar.andi %valid, %padding : i1 + scf.if %write { + %slot = index.assume %slot0 [lt(%slot0, %capacity)] : index + view.store %zero_i32, %selection[%row, %slot] : i32, view<[%nq]x[%capacity]xi32> + view.store %negative_infinity, %target[%row, %slot] : f16, view<[%nq]x[%capacity]xf16> + } + } + kernel.return +} + +kernel.def target(@ggml_compact_mask_wave64) @ggml_gather_sparse_kv_f16(%query_count: index, %key_count: index) { + %capacity = config.get @ggml.sparse_attention.capacity : index + %one = index.constant 1 : index + %four = index.constant 4 : index + %threads = index.constant 256 : index + %groups = index.div %capacity, %four : index + kernel.launch.config workgroups(%groups, %query_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%query_count: index, %key_count: index, %indices: buffer, %key: buffer, %compact_key: buffer) { + %nq = index.assume %query_count [range(%query_count, 1, 2048)] : index + %nk = index.assume %key_count [range(%key_count, 1, 131072)] : index + %capacity = config.get @ggml.sparse_attention.capacity : index + %query0 = kernel.workgroup.id : index + %query = index.assume %query0 [lt(%query0, %nq)] : index + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %four = index.constant 4 : index + %eight = index.constant 8 : index + %sixty_four = index.constant 64 : index + %offset = index.constant 0 : offset + %local_row = index.div %thread, %sixty_four : index + %local_column = index.rem %thread, %sixty_four : index + %channel = index.mul %local_column, %eight : index + %row0 = index.madd %group, %four, %local_row : index + %row = index.assume %row0 [lt(%row0, %capacity)] : index + %selection = buffer.view %indices[%offset] : buffer -> view<[%nq]x[%capacity]xi32> + %index_i32 = view.load %selection[%query, %row] : view<[%nq]x[%capacity]xi32> -> i32 + %index0 = index.cast %index_i32 : i32 to index + %index = index.assume %index0 [range(%index0, 0, 131071), lt(%index0, %nk)] : index + %source = buffer.view %key[%offset] : buffer -> view<[%nk]x512xf16> + %target = buffer.view %compact_key[%offset] : buffer -> view<[%nq]x[%capacity]x512xf16> + %values = vector.load %source[%index, %channel] : view<[%nk]x512xf16> -> vector<8xf16> + vector.store %values, %target[%query, %row, %channel] : vector<8xf16>, view<[%nq]x[%capacity]x512xf16> + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/lightning_indexer_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/lightning_indexer_f32_f16.loom new file mode 100644 index 000000000000..de3da911ec29 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/lightning_indexer_f32_f16.loom @@ -0,0 +1,164 @@ +// Reuse each query across sixteen cache keys and each key across the query heads. +amdgpu.target @ggml_indexer_wave64 {subgroup_size = 64} + +config.decl @ggml.indexer.head_size : %value: index where [range(%value, 16, 256), multiple_of(%value, 16)] + +config.decl @ggml.indexer.head_count : %value: index where [range(%value, 1, 64)] + +kernel.def target(@ggml_indexer_wave64) @ggml_lightning_indexer_f32_f16(%query_count: index, %key_count: index) { + %one = index.constant 1 : index + %sixteen = index.constant 16 : index + %fifteen = index.constant 15 : index + %threads = index.constant 256 : index + %rounded = index.add %key_count, %fifteen : index + %groups = index.div %rounded, %sixteen : index + kernel.launch.config workgroups(%groups, %query_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%query_count: index, %key_count: index, %query: buffer, %key: buffer, %weight: buffer, %mask: buffer, %output: buffer) { + %nq = index.assume %query_count [range(%query_count, 1, 2048)] : index + %nk = index.assume %key_count [range(%key_count, 1, 131072)] : index + %query_id0 = kernel.workgroup.id : index + %query_id = index.assume %query_id0 [lt(%query_id0, %nq)] : index + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : 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 + %head_size = config.get @ggml.indexer.head_size : index + %head_count = config.get @ggml.indexer.head_count : index + %packets_per_head = index.div %head_size, %c4 : index + %c256 = index.constant 256 : index + %c15 = index.constant 15 : index + %rounded_heads = index.add %head_count, %c15 : index + %head_tiles = index.div %rounded_heads, %c16 : index + %padded_heads = index.mul %head_tiles, %c16 : index + %stride = index.add %head_size, %c8 : index + %query_elements = index.mul %padded_heads, %head_size : index + %key_elements = index.mul %c16, %head_size : index + %c1023 = index.constant 1023 : index + %query_rounded = index.add %query_elements, %c1023 : index + %key_rounded = index.add %key_elements, %c1023 : index + %c1024 = index.constant 1024 : index + %query_loads = index.div %query_rounded, %c1024 : index + %key_loads = index.div %key_rounded, %c1024 : index + %zero_offset = index.constant 0 : offset + %query_stage_elements = index.mul %padded_heads, %stride : index + %key_stage_elements = index.mul %c16, %stride : index + %score_stage_elements = index.mul %c16, %padded_heads : index + %query_size = index.mul %query_stage_elements, %c2 : index + %key_size = index.mul %key_stage_elements, %c2 : index + %score_size = index.mul %score_stage_elements, %c4 : index + %query_bytes = index.cast %query_size : index to offset + %key_bytes = index.cast %key_size : index to offset + %score_bytes = index.cast %score_size : index to offset + %zero = scalar.constant 0.0 : f32 + %zero_half = vector.constant 0.0 : vector<4xf16> + %query_stage = buffer.alloca align(16) %query_bytes : buffer + %key_stage = buffer.alloca align(16) %key_bytes : buffer + %score_stage = buffer.alloca align(16) %score_bytes : buffer + %padded = encoding.layout.strided [%stride, 1] : encoding + %transposed = encoding.layout.strided [1, %stride] : encoding + %qv = buffer.view %query[%zero_offset] : buffer -> view<[%nq]x[%head_count]x[%head_size]xf32> + %kv = buffer.view %key[%zero_offset] : buffer -> view<[%nk]x[%head_size]xf16> + %wv = buffer.view %weight[%zero_offset] : buffer -> view<[%nq]x[%head_count]xf32> + %mv = buffer.view %mask[%zero_offset] : buffer -> view<[%nq]x[%nk]xf16> + %ov = buffer.view %output[%zero_offset] : buffer -> view<[%nq]x[%nk]xf32> + %qs = buffer.view %query_stage[%zero_offset] : buffer -> view<[%padded_heads]x[%head_size]xf16, %padded> + %qt = buffer.view %query_stage[%zero_offset] : buffer -> view<[%head_size]x[%padded_heads]xf16, %transposed> + %ks = buffer.view %key_stage[%zero_offset] : buffer -> view<16x[%head_size]xf16, %padded> + %scores = buffer.view %score_stage[%zero_offset] : buffer -> view<16x[%padded_heads]xf32> + %key_origin = index.mul %group, %c16 : index + scf.for %load = [%c0 to %query_loads step %c1] unroll { + %packet = index.madd %load, %c256, %thread : index + %head = index.div %packet, %packets_per_head : index + %in_tile = index.cmp ult, %head, %padded_heads : index + scf.if %in_tile { + %channel_packet = index.rem %packet, %packets_per_head : index + %channel = index.mul %channel_packet, %c4 : index + %valid_head = index.cmp ult, %head, %head_count : index + %half = scf.if %valid_head -> (vector<4xf16>) { + %q = vector.load %qv[%query_id, %head, %channel] : view<[%nq]x[%head_count]x[%head_size]xf32> -> vector<4xf32> + %value = vector.fptrunc %q : vector<4xf32> to vector<4xf16> + scf.yield %value : vector<4xf16> + } else { + scf.yield %zero_half : vector<4xf16> + } + vector.store %half, %qs[%head, %channel] : vector<4xf16>, view<[%padded_heads]x[%head_size]xf16, %padded> + } + } + scf.for %load = [%c0 to %key_loads step %c1] unroll { + %packet = index.madd %load, %c256, %thread : index + %row = index.div %packet, %packets_per_head : index + %in_tile = index.cmp ult, %row, %c16 : index + scf.if %in_tile { + %channel_packet = index.rem %packet, %packets_per_head : index + %channel = index.mul %channel_packet, %c4 : index + %key_id = index.add %key_origin, %row : index + %valid = index.cmp ult, %key_id, %nk : index + %k = scf.if %valid -> (vector<4xf16>) { + %bounded = index.assume %key_id [lt(%key_id, %nk)] : index + %value = vector.load %kv[%bounded, %channel] : view<[%nk]x[%head_size]xf16> -> vector<4xf16> + scf.yield %value : vector<4xf16> + } else { + scf.yield %zero_half : vector<4xf16> + } + vector.store %k, %ks[%row, %channel] : vector<4xf16>, view<16x[%head_size]xf16, %padded> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %active_subgroup = index.cmp ult, %subgroup, %head_tiles : index + scf.if %active_subgroup { + %head_origin = index.mul %subgroup, %c16 : index + %init_values = vector.constant 0.0 : vector<4xf32> + %init = vector.fragment %init_values shape [%c16, %c16] : vector<4xf32> + %product = scf.for %channel = [%c0 to %head_size step %c16](%acc = %init : vector<4xf32>) -> (vector<4xf32>) unroll schedule(recurrence) { + %k = vector.fragment.load %ks[%c0, %channel] shape [%c16, %c16] : view<16x[%head_size]xf16, %padded> -> vector<16xf16> + %q = vector.fragment.load %qt[%channel, %head_origin] shape [%c16, %c16] : view<[%head_size]x[%padded_heads]xf16, %transposed> -> vector<16xf16> + %next = vector.mma %k, %q, %acc : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next : vector<4xf32> + } + vector.fragment.store %product, %scores[%c0, %head_origin] shape [%c16, %c16] : vector<4xf32>, view<16x[%padded_heads]xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row = index.div %thread, %c16 : index + %head_lane = index.rem %thread, %c16 : index + %partial = scf.for %part = [%c0 to %head_tiles step %c1](%acc = %zero : f32) -> (f32) unroll { + %head = index.madd %part, %c16, %head_lane : index + %valid_head = index.cmp ult, %head, %head_count : index + %next = scf.if %valid_head -> (f32) { + %score = view.load %scores[%row, %head] : view<16x[%padded_heads]xf32> -> f32 + %relu = scalar.maxnumf %score, %zero : f32 + %w = view.load %wv[%query_id, %head] : view<[%nq]x[%head_count]xf32> -> f32 + %value = scalar.fmaf %relu, %w, %acc : f32 + scf.yield %value : f32 + } else { + scf.yield %acc : f32 + } + scf.yield %next : f32 + } + %one_i32 = scalar.constant 1 : i32 + %width = scalar.constant 16 : i32 + %total = scf.for %stage = [%c0 to %c4 step %c1](%acc = %partial : f32) -> (f32) unroll { + %shift = index.cast %stage : index to i32 + %distance = scalar.shli %one_i32, %shift : i32 + %peer, %valid = kernel.subgroup.shuffle %acc, %distance, %width : f32, i32, i32 + %next = scalar.addf %acc, %peer : f32 + scf.yield %next : f32 + } + %key_id = index.add %key_origin, %row : index + %leader = index.cmp eq, %head_lane, %c0 : index + %in_bounds = index.cmp ult, %key_id, %nk : index + %write = scalar.andi %leader, %in_bounds : i1 + scf.if %write { + %bounded = index.assume %key_id [lt(%key_id, %nk)] : index + %mask_half = view.load %mv[%query_id, %bounded] : view<[%nq]x[%nk]xf16> -> f16 + %mask_value = scalar.extf %mask_half : f16 to f32 + %result = scalar.addf %total, %mask_value : f32 + view.store %result, %ov[%query_id, %bounded] : f32, view<[%nq]x[%nk]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom index aac1e137b268..5c465e037205 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom @@ -38,21 +38,21 @@ config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range // Bounds the context copied by the test-only row-extraction kernel below. config.decl @ggml.flash_attention.test.context_capacity : %value: index where [range(%value, 1, 32768)] -template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) +template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer, %query_offset: offset) -template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) +template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer, %key_offset: offset) template.decl @ggml.flash_attention.qk_block_wmma.prefill_tail(%qk_head_size: index, %query_stage_stride: index, %tail_stage_head_size: index, %tail_score_wave: i1, %subgroup: index, %tail_key_value_stage: buffer, %query_stage: buffer, %score_stage: buffer) template.decl @ggml.flash_attention.online_softmax.prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) -template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer, %value_offset: offset) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) -template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer, %output_offset: offset, %f16_output_offset: offset) -template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %initial_max: f32, %initial_sum: f32, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer, %query_tile: index, %query_head0: index, %output_group0: index, %query_offset: offset, %key_offset: offset, %value_offset: offset, %mask_offset: offset, %output_offset: offset, %f16_output_offset: offset) kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) { %query_head_count = config.get @ggml.flash_attention.query_head_count : index @@ -73,7 +73,57 @@ kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_ kernel.launch.config workgroups(%query_tile_count, %query_head_count, %output_group_count) workgroup_size(%c256, %c1, %c1) : index } launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) where [range(%query_token_count, 1, 2048)] { %publish_f16 = scalar.constant false : i1 - template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %output) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + %initial_max = scalar.constant -1e+30 : f32 + %initial_sum = scalar.constant 0.0 : f32 + %query_tile = kernel.workgroup.id : index + %query_head0 = kernel.workgroup.id : index + %output_group0 = kernel.workgroup.id : index + %query_offset = index.constant 0 : offset + %key_offset = index.constant 0 : offset + %value_offset = index.constant 0 : offset + %mask_offset = index.constant 0 : offset + %output_offset = index.constant 0 : offset + %f16_output_offset = index.constant 0 : offset + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %initial_max, %initial_sum, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %output, %query_tile, %query_head0, %output_group0, %query_offset, %key_offset, %value_offset, %mask_offset, %output_offset, %f16_output_offset) : (i1, f32, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, index, index, index, offset, offset, offset, offset, offset, offset) + kernel.return +} + +kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma_sink(%query_token_count: index, %key_value_token_count: index) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), multiple_of(%value_head_size0, 64)] : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_query_token_count = index.add %query_token_count, %c15 : index + %query_tile_count = index.div %padded_query_token_count, %c16 : index + %output_tile_count = index.div %value_head_size, %c64 : index + %padded_output_tile_count = index.add %output_tile_count, %c3 : index + %output_group_count = index.div %padded_output_tile_count, %c4 : index + kernel.launch.config workgroups(%query_tile_count, %query_head_count, %output_group_count) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %sink: buffer, %output: buffer) where [range(%query_token_count, 1, 2048)] { + %publish_f16 = scalar.constant false : i1 + %head_count = config.get @ggml.flash_attention.query_head_count : index + %head0 = kernel.workgroup.id : index + %head, %bounded_head_count = index.assume %head0, %head_count [lt(%head0, %head_count)] : index, index + %zero_offset = index.constant 0 : offset + %sink_view = buffer.view %sink[%zero_offset] : buffer -> view<[%bounded_head_count]xf32> + %initial_max = view.load %sink_view[%head] : view<[%bounded_head_count]xf32> -> f32 + %initial_sum = scalar.constant 1.0 : f32 + %query_tile = kernel.workgroup.id : index + %query_head0 = kernel.workgroup.id : index + %output_group0 = kernel.workgroup.id : index + %query_offset = index.constant 0 : offset + %key_offset = index.constant 0 : offset + %value_offset = index.constant 0 : offset + %mask_offset = index.constant 0 : offset + %output_offset = index.constant 0 : offset + %f16_output_offset = index.constant 0 : offset + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %initial_max, %initial_sum, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %sink, %output, %output, %query_tile, %query_head0, %output_group0, %query_offset, %key_offset, %value_offset, %mask_offset, %output_offset, %f16_output_offset) : (i1, f32, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, index, index, index, offset, offset, offset, offset, offset, offset) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom index 44f7c77be19b..51fd9a8c701d 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom @@ -12,7 +12,7 @@ config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value config.decl @ggml.flash_attention.gate_stride_head : %value: index where [range(%value, 1, 16777216)] config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range(%value, 1, 16777216)] -template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %initial_max: f32, %initial_sum: f32, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer, %query_tile: index, %query_head0: index, %output_group0: index, %query_offset: offset, %key_offset: offset, %value_offset: offset, %mask_offset: offset, %output_offset: offset, %f16_output_offset: offset) kernel.def target(@ggml_flash_attention_gfx11_wave64) export("ggml_flash_attention_f32_f16_wmma_publish_f16") @ggml_flash_attention_f32_f16_wmma_publish_f16(%query_token_count: index, %key_value_token_count: index) { %query_head_count = config.get @ggml.flash_attention.query_head_count : index @@ -33,6 +33,17 @@ kernel.def target(@ggml_flash_attention_gfx11_wave64) export("ggml_flash_attenti kernel.launch.config workgroups(%query_tile_count, %query_head_count, %output_group_count) workgroup_size(%c256, %c1, %c1) : index } launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) where [range(%query_token_count, 1, 2048)] { %publish_f16 = scalar.constant true : i1 - template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %f16_output) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + %initial_max = scalar.constant -1e+30 : f32 + %initial_sum = scalar.constant 0.0 : f32 + %query_tile = kernel.workgroup.id : index + %query_head0 = kernel.workgroup.id : index + %output_group0 = kernel.workgroup.id : index + %query_offset = index.constant 0 : offset + %key_offset = index.constant 0 : offset + %value_offset = index.constant 0 : offset + %mask_offset = index.constant 0 : offset + %output_offset = index.constant 0 : offset + %f16_output_offset = index.constant 0 : offset + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %initial_max, %initial_sum, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %f16_output, %query_tile, %query_head0, %output_group0, %query_offset, %key_offset, %value_offset, %mask_offset, %output_offset, %f16_output_offset) : (i1, f32, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, index, index, index, offset, offset, offset, offset, offset, offset) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/sparse_prefill_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/sparse_prefill_f32_f16.loom new file mode 100644 index 000000000000..5aea0d65d92f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/sparse_prefill_f32_f16.loom @@ -0,0 +1,40 @@ +amdgpu.target @ggml_sparse_attention_wave64 {subgroup_size = 64} + +config.decl @ggml.sparse_attention.capacity : %value: index where [range(%value, 64, 2048), multiple_of(%value, 64)] + +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %initial_max: f32, %initial_sum: f32, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer, %query_tile: index, %query_head0: index, %output_group0: index, %query_offset: offset, %key_offset: offset, %value_offset: offset, %mask_offset: offset, %output_offset: offset, %f16_output_offset: offset) + +// Each query owns a compact K/V set. Adjacent workgroups reuse that set across +// four groups of sixteen heads and two output-channel groups. +kernel.def target(@ggml_sparse_attention_wave64) @ggml_sparse_prefill_f32_f16(%query_count: index) { + %one = index.constant 1 : index + %eight = index.constant 8 : index + %threads = index.constant 256 : index + kernel.launch.config workgroups(%eight, %query_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%query_count: index, %query: buffer, %key: buffer, %mask: buffer, %sink: buffer, %output: buffer) { + %nq = index.assume %query_count [range(%query_count, 1, 2048)] : index + %row0 = kernel.workgroup.id : index + %row = index.assume %row0 [lt(%row0, %nq)] : index + %group0 = kernel.workgroup.id : index + %group = index.assume %group0 [range(%group0, 0, 7)] : index + %zero = index.constant 0 : index + %four = index.constant 4 : index + %sixteen = index.constant 16 : index + %zero_offset = index.constant 0 : offset + %query_stride = index.constant 131072 : offset + %key_row_bytes = index.constant 1024 : offset + %half_bytes = index.constant 2 : offset + %capacity = config.get @ggml.sparse_attention.capacity : index + %key_stride = index.scale %capacity, %key_row_bytes : index, offset -> offset + %mask_stride = index.scale %capacity, %half_bytes : index, offset -> offset + %query_offset = index.scale %row, %query_stride : index, offset -> offset + %key_offset = index.scale %row, %key_stride : index, offset -> offset + %mask_offset = index.scale %row, %mask_stride : index, offset -> offset + %head = index.rem %group, %four : index + %output_group = index.div %group, %four : index + %publish_f16 = scalar.constant false : i1 + %initial_max = scalar.constant 0.0 : f32 + %initial_sum = scalar.constant 1.0 : f32 + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %initial_max, %initial_sum, %sixteen, %capacity, %query, %key, %key, %mask, %sink, %output, %output, %zero, %head, %output_group, %query_offset, %key_offset, %key_offset, %mask_offset, %query_offset, %zero_offset) : (i1, f32, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, index, index, index, offset, offset, offset, offset, offset, offset) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/hadamard_128_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/hadamard_128_f32.loom new file mode 100644 index 000000000000..a643a69ee410 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/hadamard_128_f32.loom @@ -0,0 +1,81 @@ +// Normalized Sylvester transform, using the caller's Hadamard semantic hint. +// Four rows per workgroup; one wave keeps all 128 row values in registers. +amdgpu.target @ggml_hadamard_wave32 {subgroup_size = 32} + +func.def inline @ggml_hadamard_butterfly(%value: f32, %lane: i32, %mask: i32) -> (f32) { + %width = scalar.constant 32 : i32 + %zero = scalar.constant 0 : i32 + %peer, %valid = kernel.subgroup.shuffle %value, %mask, %width : f32, i32, i32 + %bit = scalar.andi %lane, %mask : i32 + %left = scalar.cmpi eq, %bit, %zero : i32 + %sum = scalar.addf %value, %peer : f32 + %difference = scalar.subf %peer, %value : f32 + %result = scf.select %left, %sum, %difference : f32 + func.return %result : f32 +} + +kernel.def target(@ggml_hadamard_wave32) @ggml_hadamard_128_f32(%row_count: index) { + %one = index.constant 1 : index + %four = index.constant 4 : index + %three = index.constant 3 : index + %size = index.constant 128 : index + %padded = index.add %row_count, %three : index + %groups = index.div %padded, %four : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%size, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %rows = index.assume %row_count [range(%row_count, 1, 1048576)] : index + %group = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %zero_offset = index.constant 0 : offset + %one_bit = scalar.constant 1 : i32 + %scale = scalar.constant 0.088388347648318447 : f32 + %row0 = index.madd %group, %c4, %subgroup : index + %active = index.cmp ult, %row0, %rows : index + scf.if %active { + %row = index.assume %row0 [lt(%row0, %rows)] : index + %i1 = index.add %lane, %c32 : index + %i2 = index.add %lane, %c64 : index + %i3 = index.add %lane, %c96 : index + %lane_i32 = index.cast %lane : index to i32 + %x = buffer.view %input[%zero_offset] : buffer -> view<[%rows]x128xf32> + %y = buffer.view %output[%zero_offset] : buffer -> view<[%rows]x128xf32> + %x0 = view.load %x[%row, %lane] : view<[%rows]x128xf32> -> f32 + %x1 = view.load %x[%row, %i1] : view<[%rows]x128xf32> -> f32 + %x2 = view.load %x[%row, %i2] : view<[%rows]x128xf32> -> f32 + %x3 = view.load %x[%row, %i3] : view<[%rows]x128xf32> -> f32 + %r0, %r1, %r2, %r3 = scf.for %stage = [%c0 to %c5 step %c1](%v0 = %x0 : f32, %v1 = %x1 : f32, %v2 = %x2 : f32, %v3 = %x3 : f32) -> (f32, f32, f32, f32) unroll { + %shift = index.cast %stage : index to i32 + %mask = scalar.shli %one_bit, %shift : i32 + %n0 = func.call @ggml_hadamard_butterfly(%v0, %lane_i32, %mask) : (f32, i32, i32) -> (f32) + %n1 = func.call @ggml_hadamard_butterfly(%v1, %lane_i32, %mask) : (f32, i32, i32) -> (f32) + %n2 = func.call @ggml_hadamard_butterfly(%v2, %lane_i32, %mask) : (f32, i32, i32) -> (f32) + %n3 = func.call @ggml_hadamard_butterfly(%v3, %lane_i32, %mask) : (f32, i32, i32) -> (f32) + scf.yield %n0, %n1, %n2, %n3 : f32, f32, f32, f32 + } + %a0 = scalar.addf %r0, %r1 : f32 + %a1 = scalar.subf %r0, %r1 : f32 + %a2 = scalar.addf %r2, %r3 : f32 + %a3 = scalar.subf %r2, %r3 : f32 + %b0 = scalar.addf %a0, %a2 : f32 + %b1 = scalar.addf %a1, %a3 : f32 + %b2 = scalar.subf %a0, %a2 : f32 + %b3 = scalar.subf %a1, %a3 : f32 + %y0 = scalar.mulf %b0, %scale : f32 + %y1 = scalar.mulf %b1, %scale : f32 + %y2 = scalar.mulf %b2, %scale : f32 + %y3 = scalar.mulf %b3, %scale : f32 + view.store %y0, %y[%row, %lane] : f32, view<[%rows]x128xf32> + view.store %y1, %y[%row, %i1] : f32, view<[%rows]x128xf32> + view.store %y2, %y[%row, %i2] : f32, view<[%rows]x128xf32> + view.store %y3, %y[%row, %i3] : f32, view<[%rows]x128xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/small_k_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/small_k_f32.loom new file mode 100644 index 000000000000..0e607d05d741 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/small_k_f32.loom @@ -0,0 +1,97 @@ +config.decl @ggml.small_matmul.k : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.m : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.n : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.repeat2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.repeat3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.weight_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.input_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.weight_format : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.ws0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.ws1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.ws2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.ws3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.xs0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.xs1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.xs2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.small_matmul.xs3 : %value: index where [range(%value, 1, 1073741824)] + +kernel.def @ggml_mul_mat_small_k_f32(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %weight: buffer, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %k = config.get @ggml.small_matmul.k : index + %m = config.get @ggml.small_matmul.m : index + %n = config.get @ggml.small_matmul.n : index + %ne2 = config.get @ggml.small_matmul.ne2 : index + %repeat2 = config.get @ggml.small_matmul.repeat2 : index + %repeat3 = config.get @ggml.small_matmul.repeat3 : index + %weight_span = config.get @ggml.small_matmul.weight_span : index + %input_span = config.get @ggml.small_matmul.input_span : index + %weight_format = config.get @ggml.small_matmul.weight_format : index + %ws0 = config.get @ggml.small_matmul.ws0 : index + %ws1 = config.get @ggml.small_matmul.ws1 : index + %ws2 = config.get @ggml.small_matmul.ws2 : index + %ws3 = config.get @ggml.small_matmul.ws3 : index + %xs0 = config.get @ggml.small_matmul.xs0 : index + %xs1 = config.get @ggml.small_matmul.xs1 : index + %xs2 = config.get @ggml.small_matmul.xs2 : index + %xs3 = config.get @ggml.small_matmul.xs3 : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear0, %count : index + scf.if %valid { + %linear, %extent = index.assume %linear0, %count [range(%linear0, 0, 134217727), lt(%linear0, %count)] : index, index + %row = index.rem %linear, %m : index + %rest0 = index.div %linear, %m : index + %column = index.rem %rest0, %n : index + %rest1 = index.div %rest0, %n : index + %batch2 = index.rem %rest1, %ne2 : index + %batch3 = index.div %rest1, %ne2 : index + %wb2 = index.div %batch2, %repeat2 : index + %wb3 = index.div %batch3, %repeat3 : index + %wbase0 = index.mul %row, %ws1 : index + %wbase1 = index.madd %wb2, %ws2, %wbase0 : index + %wbase = index.madd %wb3, %ws3, %wbase1 : index + %xbase0 = index.mul %column, %xs1 : index + %xbase1 = index.madd %batch2, %xs2, %xbase0 : index + %xbase = index.madd %batch3, %xs3, %xbase1 : index + %zero = index.constant 0 : index + %one = index.constant 1 : index + %zero_offset = index.constant 0 : offset + %zero_f32 = scalar.constant 0.0 : f32 + %format16 = index.constant 16 : index + %is_f16 = index.cmp eq, %weight_format, %format16 : index + %wv16 = buffer.view %weight[%zero_offset] : buffer -> view<[%weight_span]xf16> + %wv32 = buffer.view %weight[%zero_offset] : buffer -> view<[%weight_span]xf32> + %xv = buffer.view %input[%zero_offset] : buffer -> view<[%input_span]xf32> + %yv = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + %sum = scf.for %column_k = [%zero to %k step %one](%acc = %zero_f32 : f32) -> (f32) { + %wi0 = index.madd %column_k, %ws0, %wbase : index + %xi0 = index.madd %column_k, %xs0, %xbase : index + %wi, %we = index.assume %wi0, %weight_span [range(%wi0, 0, 1073741823), lt(%wi0, %weight_span)] : index, index + %xi, %xe = index.assume %xi0, %input_span [range(%xi0, 0, 1073741823), lt(%xi0, %input_span)] : index, index + %w = scf.if %is_f16 -> (f32) { + %narrow = view.load %wv16[%wi] : view<[%weight_span]xf16> -> f16 + %wide = scalar.extf %narrow : f16 to f32 + scf.yield %wide : f32 + } else { + %wide = view.load %wv32[%wi] : view<[%weight_span]xf32> -> f32 + scf.yield %wide : f32 + } + %x = view.load %xv[%xi] : view<[%input_span]xf32> -> f32 + %product = scalar.mulf %w, %x : f32 + %next = scalar.addf %acc, %product : f32 + scf.yield %next : f32 + } + view.store %sum, %yv[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q8_q8_1.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q8_q8_1.loom new file mode 100644 index 000000000000..107d084d7505 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q8_q8_1.loom @@ -0,0 +1,96 @@ +// One wave per output row. Shared Q8_1_x4 activations are produced once +// and reused across rows/projections; raw Q8_0 weights stay resident. +amdgpu.target @ggml_q8_dot_wave32 {subgroup_size = 32} + +config.decl @ggml.matmul.vector.input_size : %v: index where [range(%v, 256, 32768), multiple_of(%v, 256)] + +config.decl @ggml.matmul.vector.output_size : %v: index where [range(%v, 1, 524288)] + +config.def @ggml.matmul.vector.batch_count = 1 : index + +kernel.def target(@ggml_q8_dot_wave32) @ggml_mul_mat_vector_q8_q8_1(%token_count: index) { + %m = config.get @ggml.matmul.vector.output_size : index + %batches = config.get @ggml.matmul.vector.batch_count : index + %one = index.constant 1 : index + %wave = index.constant 32 : index + kernel.launch.config workgroups(%m, %batches, %one) workgroup_size(%wave, %one, %one) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %next_output: buffer) { + %k = config.get @ggml.matmul.vector.input_size : index + %m = config.get @ggml.matmul.vector.output_size : index + %batch_count = config.get @ggml.matmul.vector.batch_count : index + %batches = index.assume %batch_count [range(%batch_count, 1, 512)] : index + %batch0 = kernel.workgroup.id : index + %batch = index.assume %batch0 [lt(%batch0, %batches)] : index + %row0 = kernel.workgroup.id : index + %row, %rows = index.assume %row0, %m [lt(%row0, %m)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %c128 = index.constant 128 : index + %c32 = index.constant 32 : index + %c512 = index.constant 512 : index + %zero_offset = index.constant 0 : offset + %two_bytes = index.constant 2 : offset + %block_bytes = index.constant 34 : offset + %activation_group_bytes = index.constant 144 : offset + %activation_header_bytes = index.constant 16 : offset + %integer_zero = vector.constant 0 : vector<4xi32> + %scalar_zero = scalar.constant 0 : i32 + %zero_float = scalar.constant 0.0 : f32 + %weight_blocks = index.div %k, %c32 : index + %row_bytes = index.scale %weight_blocks, %block_bytes : index, offset -> offset + %flat_row = index.madd %batch, %rows, %row : index + %flat_rows = index.mul %batches, %rows : index + %row_base = index.scale %flat_row, %row_bytes : index, offset -> offset + %activation_groups = index.div %k, %c128 : index + %activation_row_bytes = index.scale %activation_groups, %activation_group_bytes : index, offset -> offset + %activation_row_base = index.scale %batch, %activation_row_bytes : index, offset -> offset + %lane_k = index.mul %lane, %c16 : index + %packet = index.rem %lane_k, %c32 : index + %lane_sum = scf.for %base = [%c0 to %k step %c512](%acc = %zero_float : f32) -> (f32) { + %ki0 = index.add %base, %lane_k : index + %active = index.cmp ult, %ki0, %k : index + %next = scf.if %active -> (f32) { + %ki = index.assume %ki0 [lt(%ki0, %k), multiple_of(%ki0, 16)] : index + %block = index.div %ki, %c32 : index + %block_offset = index.scale %block, %block_bytes : index, offset -> offset + %scale_offset = index.add %row_base, %block_offset : offset + %payload_offset = index.add %scale_offset, %two_bytes : offset + %sv = buffer.view %weight[%scale_offset] : buffer -> view<1xf16> + %wv = buffer.view %weight[%payload_offset] : buffer -> view<32xi8> + %half_scale = view.load %sv[%c0] : view<1xf16> -> f16 + %scale = scalar.extf %half_scale : f16 to f32 + %packed = vector.load %wv[%packet] : view<32xi8> -> vector<16xi8> + %activation_group = index.div %ki, %c128 : index + %activation_group_base = index.scale %activation_group, %activation_group_bytes : index, offset -> offset + %activation_base = index.add %activation_row_base, %activation_group_base : offset + %activation_payload = index.add %activation_base, %activation_header_bytes : offset + %activation_index = index.rem %ki, %c128 : index + %activation_block = index.rem %block, %c4 : index + %ds_index = index.mul %activation_block, %c2 : index + %xv = buffer.view %input[%activation_payload] : buffer -> view<128xi8> + %ds = buffer.view %input[%activation_base] : buffer -> view<8xf16> + %a = vector.load %xv[%activation_index] : view<128xi8> -> vector<16xi8> + %input_half_scale = view.load %ds[%ds_index] : view<8xf16> -> f16 + %input_scale = scalar.extf %input_half_scale : f16 to f32 + %dot_pair = vector.dot4i %packed, %a, %integer_zero : vector<16xi8>, vector<16xi8>, vector<4xi32> + %dot_integer = vector.reduce %dot_pair, %scalar_zero : vector<4xi32>, i32 + %dot = scalar.sitofp %dot_integer : i32 to f32 + %product_scale = scalar.mulf %scale, %input_scale : f32 + %value = scalar.fmaf %product_scale, %dot, %acc : f32 + scf.yield %value : f32 + } else { + scf.yield %acc : f32 + } + scf.yield %next : f32 + } + %total = kernel.subgroup.reduce %lane_sum : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %y = buffer.view %output[%zero_offset] : buffer -> view<[%flat_rows]xf32> + view.store %total, %y[%flat_row] : f32, view<[%flat_rows]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_small_m_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_small_m_f32.loom new file mode 100644 index 000000000000..b165a6eaabc4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_small_m_f32.loom @@ -0,0 +1,62 @@ +// One workgroup per row keeps long reductions parallel for narrow projections. +amdgpu.target @ggml_small_m_wave32 {subgroup_size = 32} +config.decl @ggml.small_m.k : %v: index where [range(%v, 4096, 32768), multiple_of(%v, 2048)] +config.decl @ggml.small_m.m : %v: index where [range(%v, 1, 32)] + +kernel.def target(@ggml_small_m_wave32) @ggml_mul_mat_vector_small_m_f32() { + %m = config.get @ggml.small_m.m : index + %one = index.constant 1 : index + %threads = index.constant 256 : index + kernel.launch.config workgroups(%m, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%weight: buffer, %input: buffer, %output: buffer) { + %k = config.get @ggml.small_m.k : index + %m = config.get @ggml.small_m.m : index + %row0 = kernel.workgroup.id : index + %row = index.assume %row0 [lt(%row0, %m)] : index + %thread = kernel.workitem.id : index + %lane = kernel.subgroup.lane.id : index + %wave0 = kernel.subgroup.id : index + %wave = index.assume %wave0 [range(%wave0, 0, 7)] : index + %zero = index.constant 0 : index + %eight = index.constant 8 : index + %step = index.constant 2048 : index + %zero_offset = index.constant 0 : offset + %shared_bytes = index.constant 32 : offset + %partial = buffer.alloca align(16) %shared_bytes : buffer + %pv = buffer.view %partial[%zero_offset] : buffer -> view<8xf32> + %wv = buffer.view %weight[%zero_offset] : buffer -> view<[%m]x[%k]xf32> + %xv = buffer.view %input[%zero_offset] : buffer -> view<[%k]xf32> + %yv = buffer.view %output[%zero_offset] : buffer -> view<[%m]xf32> + %thread_k = index.mul %thread, %eight : index + %zeros = vector.constant 0.0 : vector<8xf32> + %acc = scf.for %base = [%zero to %k step %step](%sum = %zeros : vector<8xf32>) -> (vector<8xf32>) { + %col = index.add %base, %thread_k : index + %w = vector.load %wv[%row, %col] : view<[%m]x[%k]xf32> -> vector<8xf32> + %x = vector.load %xv[%col] : view<[%k]xf32> -> vector<8xf32> + %next = vector.fmaf %w, %x, %sum : vector<8xf32> + scf.yield %next : vector<8xf32> + } + %zero_float = scalar.constant 0.0 : f32 + %local = vector.reduce %acc, %zero_float : vector<8xf32>, f32 + %wave_sum = kernel.subgroup.reduce %local : f32 + %leader = index.cmp eq, %lane, %zero : index + scf.if %leader { + view.store %wave_sum, %pv[%wave] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %first_wave = index.cmp eq, %wave, %zero : index + scf.if %first_wave { + %has_partial = index.cmp ult, %lane, %eight : index + %part = scf.if %has_partial -> (f32) { + %v = view.load %pv[%lane] : view<8xf32> -> f32 + scf.yield %v : f32 + } else { + scf.yield %zero_float : f32 + } + %total = kernel.subgroup.reduce %part : f32 + scf.if %leader { + view.store %total, %yv[%row] : f32, view<[%m]xf32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mxfp4_prefill_prefetch.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mxfp4_prefill_prefetch.loom new file mode 100644 index 000000000000..e4c2839b96b4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mxfp4_prefill_prefetch.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.mxfp4_prefetch.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.mxfp4_prefetch.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.mxfp4_prefetch.publish_vector8(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<8xf32>, %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_mxfp4_prefill_prefetch_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), multiple_of(%value, 256)] + +config.decl @ggml.mul_mat_id.output_size : %value: index where [range(%value, 64, 262144), multiple_of(%value, 64)] + +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.mxfp4_prefetch.publish_vector8> device @ggml_mul_mat_id_mxfp4_prefetch_publish_vector8(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<8xf32>, %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<8xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<8xf32>, view<[%output_row_count]x[%output_size]xf32>, vector<8xi1> + } + template.return +} + +template.def<@ggml.matmul_id.mxfp4_prefetch.finish_tile> device @ggml_mul_mat_id_mxfp4_prefetch_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_mxfp4_prefill_prefetch_gfx11_wave32) @ggml_mul_mat_id_mxfp4_prefill_prefetch(%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(%c64, %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.mxfp4_prefetch.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/vector_mxfp4_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_mxfp4_f32.loom new file mode 100644 index 000000000000..f132ed8a44af --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_mxfp4_f32.loom @@ -0,0 +1,121 @@ +// One wave owns one output row and one route. Each lane decodes eight values +// from four GGUF bytes; activations are reused from cache across output rows. +amdgpu.target @ggml_mxfp4_vector_wave32 {subgroup_size = 32} + +config.decl @ggml.mxfp4_vector.input_size : %v: index where [range(%v, 256, 32768), multiple_of(%v, 256)] + +config.decl @ggml.mxfp4_vector.output_size : %v: index where [range(%v, 1, 262144)] + +config.decl @ggml.mxfp4_vector.route_count : %v: index where [range(%v, 1, 32)] + +config.decl @ggml.mxfp4_vector.input_route_count : %v: index where [range(%v, 1, 32)] + +config.decl @ggml.mxfp4_vector.expert_count : %v: index where [range(%v, 1, 512)] + +kernel.def target(@ggml_mxfp4_vector_wave32) @ggml_mul_mat_id_vector_mxfp4_f32(%token_count: index) { + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %one = index.constant 1 : index + %wave = index.constant 32 : index + kernel.launch.config workgroups(%m, %routes, %one) workgroup_size(%wave, %one, %one) : index +} launch(%token_count: index, %input: buffer, %ids: buffer, %weight: buffer, %output: buffer) { + %k = config.get @ggml.mxfp4_vector.input_size : index + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %input_routes = config.get @ggml.mxfp4_vector.input_route_count : index + %experts = config.get @ggml.mxfp4_vector.expert_count : index + %row0 = kernel.workgroup.id : index + %route0 = kernel.workgroup.id : index + %row, %rows = index.assume %row0, %m [lt(%row0, %m)] : index, index + %route, %route_extent = index.assume %route0, %routes [lt(%route0, %routes)] : index, index + %lane = kernel.subgroup.lane.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 + %c5 = index.constant 5 : index + %c6 = index.constant 6 : index + %c7 = index.constant 7 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %zero_offset = index.constant 0 : offset + %one_byte = index.constant 1 : offset + %block_bytes = index.constant 17 : offset + %zero_float = scalar.constant 0.0 : f32 + %zero_i32 = scalar.constant 0 : i32 + %scale_shift = scalar.constant 23 : i32 + %scale_subnormal = scalar.constant 4194304 : i32 + %scale_nan = scalar.constant 2143289344 : i32 + %nan_exponent = scalar.constant 255 : i32 + %id_view = buffer.view %ids[%zero_offset] : buffer -> view<[%route_extent]xi32> + %expert_i32 = view.load %id_view[%route] : view<[%route_extent]xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %valid = index.cmp ult, %expert0, %experts : index + %sum = scf.if %valid -> (f32) { + %expert, %expert_extent = index.assume %expert0, %experts [lt(%expert0, %experts)] : index, index + %weight_blocks = index.div %k, %c32 : index + %row_bytes = index.scale %weight_blocks, %block_bytes : index, offset -> offset + %flat_row = index.madd %expert, %rows, %row : index + %row_base = index.scale %flat_row, %row_bytes : index, offset -> offset + %input_route = index.rem %route, %input_routes : index + %x = buffer.view %input[%zero_offset] : buffer -> view<[%input_routes]x[%k]xf32> + %lane_block = index.div %lane, %c4 : index + %lane_packet = index.rem %lane, %c4 : index + %packet_offset = index.mul %lane_packet, %c4 : index + %lane_k = index.madd %lane_block, %c32, %packet_offset : index + %lane_sum = scf.for %base = [%c0 to %k step %c256](%acc = %zero_float : f32) -> (f32) { + %klo0 = index.add %base, %lane_k : index + %khi0 = index.add %klo0, %c16 : index + %klo = index.assume %klo0 [lt(%klo0, %k), multiple_of(%klo0, 4)] : index + %khi = index.assume %khi0 [lt(%khi0, %k), multiple_of(%khi0, 4)] : index + %block = index.div %klo, %c32 : index + %block_offset = index.scale %block, %block_bytes : index, offset -> offset + %scale_offset = index.add %row_base, %block_offset : offset + %payload_offset = index.add %scale_offset, %one_byte : offset + %scv = buffer.view %weight[%scale_offset] : buffer -> view<1xi8> + %wv = buffer.view %weight[%payload_offset] : buffer -> view<16xi8> + %scale_byte = view.load %scv[%c0] : view<1xi8> -> i8 + %exponent = scalar.extui %scale_byte : i8 to i32 + %normal_bits = scalar.shli %exponent, %scale_shift : i32 + %is_zero = scalar.cmpi eq, %exponent, %zero_i32 : i32 + %is_nan = scalar.cmpi eq, %exponent, %nan_exponent : i32 + %finite_bits = scf.select %is_zero, %scale_subnormal, %normal_bits : i32 + %scale_bits = scf.select %is_nan, %scale_nan, %finite_bits : i32 + %scale = scalar.bitcast %scale_bits : i32 to f32 + %packed_bytes = vector.load %wv[%packet_offset] : view<16xi8> -> vector<4xi8> + %packed = vector.bitcast %packed_bytes : vector<4xi8> to vector<1xi32> + %schema = encoding.define #encoding.f4e2m1 : encoding + %decoded = vector.decode %packed using %schema : vector<1xi32>, encoding -> vector<8xf16> + %w0 = vector.extract %decoded[%c0] : vector<8xf16> -> f16 + %w1 = vector.extract %decoded[%c1] : vector<8xf16> -> f16 + %w2 = vector.extract %decoded[%c2] : vector<8xf16> -> f16 + %w3 = vector.extract %decoded[%c3] : vector<8xf16> -> f16 + %w4 = vector.extract %decoded[%c4] : vector<8xf16> -> f16 + %w5 = vector.extract %decoded[%c5] : vector<8xf16> -> f16 + %w6 = vector.extract %decoded[%c6] : vector<8xf16> -> f16 + %w7 = vector.extract %decoded[%c7] : vector<8xf16> -> f16 + %lo_half = vector.from_elements %w0, %w2, %w4, %w6 : vector<4xf16> + %hi_half = vector.from_elements %w1, %w3, %w5, %w7 : vector<4xf16> + %lo = vector.extf %lo_half : vector<4xf16> to vector<4xf32> + %hi = vector.extf %hi_half : vector<4xf16> to vector<4xf32> + %xlo = vector.load %x[%input_route, %klo] : view<[%input_routes]x[%k]xf32> -> vector<4xf32> + %xhi = vector.load %x[%input_route, %khi] : view<[%input_routes]x[%k]xf32> -> vector<4xf32> + %dlo = vector.dotf %lo, %xlo, %zero_float : vector<4xf32>, vector<4xf32>, f32 + %dot = vector.dotf %hi, %xhi, %dlo : vector<4xf32>, vector<4xf32>, f32 + %next = scalar.fmaf %scale, %dot, %acc : f32 + scf.yield %next : f32 + } + scf.yield %lane_sum : f32 + } else { + scf.yield %zero_float : f32 + } + %total = kernel.subgroup.reduce %sum : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %y = buffer.view %output[%zero_offset] : buffer -> view<[%route_extent]x[%rows]xf32> + view.store %total, %y[%route, %row] : f32, view<[%route_extent]x[%rows]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_mxfp4_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_mxfp4_q8.loom new file mode 100644 index 000000000000..2bea020a41c3 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_mxfp4_q8.loom @@ -0,0 +1,364 @@ +// One wave per output row/route/token; sixteen exact FP4 values per lane. +// Shared Q8_1_x4 activations are reused across expert rows and projections. +amdgpu.target @ggml_mxfp4_dot_wave32 {subgroup_size = 32} + +config.decl @ggml.mxfp4_vector.input_size : %v: index where [range(%v, 512, 32768), multiple_of(%v, 512)] + +config.decl @ggml.mxfp4_vector.output_size : %v: index where [range(%v, 1, 262144)] + +config.decl @ggml.mxfp4_vector.route_count : %v: index where [range(%v, 1, 32)] + +config.decl @ggml.mxfp4_vector.route_stride : %v: index where [range(%v, 1, 512)] + +config.decl @ggml.mxfp4_vector.input_route_count : %v: index where [range(%v, 1, 32)] + +config.decl @ggml.mxfp4_vector.expert_count : %v: index where [range(%v, 1, 512)] + +template.decl @ggml.mxfp4_vector.update(%klo: index, %block: index, %packet_offset: index, %activation_base: offset, %lo: vector<8xi8>, %hi: vector<8xi8>, %scale: f32, %acc: f32, %input: buffer) -> (f32) + +template.decl @ggml.mxfp4_vector.lane_sum(%paired: i1, %row_base: offset, %input_base: offset, %peer_base: offset, %input: buffer, %weight: buffer) -> (f32, f32) + +template.def<@ggml.mxfp4_vector.update> device @ggml_mxfp4_vector_update(%klo: index, %block: index, %packet_offset: index, %activation_base: offset, %lo: vector<8xi8>, %hi: vector<8xi8>, %scale: f32, %acc: f32, %input: buffer) -> (f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %group_bytes = index.constant 144 : offset + %header_bytes = index.constant 16 : offset + %half_float = scalar.constant 0.5 : f32 + %dot_zero = vector.constant 0 : vector<2xi32> + %zero_i32 = scalar.constant 0 : i32 + %agroup = index.div %klo, %c128 : index + %aoffset = index.scale %agroup, %group_bytes : index, offset -> offset + %abase = index.add %activation_base, %aoffset : offset + %apayload = index.add %abase, %header_bytes : offset + %ablock = index.rem %block, %c4 : index + %ai0 = index.madd %ablock, %c32, %packet_offset : index + %ai_hi0 = index.add %ai0, %c16 : index + // Four blocks per group, two eight-byte packets per nibble plane. + %ai = index.assume %ai0 [range(%ai0, 0, 104), multiple_of(%ai0, 8)] : index + %ai_hi = index.assume %ai_hi0 [range(%ai_hi0, 16, 120), multiple_of(%ai_hi0, 8)] : index + %ds_index = index.mul %ablock, %c2 : index + %av = buffer.view %input[%apayload] : buffer -> view<128xi8> + %ds = buffer.view %input[%abase] : buffer -> view<8xf16> + %xlo = vector.load %av[%ai] : view<128xi8> -> vector<8xi8> + %xhi = vector.load %av[%ai_hi] : view<128xi8> -> vector<8xi8> + %dx_half = view.load %ds[%ds_index] : view<8xf16> -> f16 + %dx = scalar.extf %dx_half : f16 to f32 + %dot_lo = vector.dot4i %lo, %xlo, %dot_zero : vector<8xi8>, vector<8xi8>, vector<2xi32> + %dot_pair = vector.dot4i %hi, %xhi, %dot_lo : vector<8xi8>, vector<8xi8>, vector<2xi32> + %dot_integer = vector.reduce %dot_pair, %zero_i32 : vector<2xi32>, i32 + %dot = scalar.sitofp %dot_integer : i32 to f32 + %fp4_scale = scalar.mulf %scale, %half_float : f32 + %product_scale = scalar.mulf %fp4_scale, %dx : f32 + %next = scalar.fmaf %product_scale, %dot, %acc : f32 + template.return %next : f32 +} + +template.def<@ggml.mxfp4_vector.lane_sum> device @ggml_mxfp4_vector_lane_sum(%paired: i1, %row_base: offset, %input_base: offset, %peer_base: offset, %input: buffer, %weight: buffer) -> (f32, f32) { + %k = config.get @ggml.mxfp4_vector.input_size : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c512 = index.constant 512 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %one_byte = index.constant 1 : offset + %block_bytes = index.constant 17 : offset + %mask4 = vector.constant 15 : vector<8xi8> + %shift4 = vector.constant 4 : vector<2xi32> + %zero_float = scalar.constant 0.0 : f32 + %zero_i32 = scalar.constant 0 : i32 + %scale_shift = scalar.constant 23 : i32 + %scale_subnormal = scalar.constant 4194304 : i32 + %scale_nan = scalar.constant 2143289344 : i32 + %nan_exponent = scalar.constant 255 : i32 + %fp4_0 = scalar.constant 0 : i8 + %fp4_1 = scalar.constant 1 : i8 + %fp4_2 = scalar.constant 2 : i8 + %fp4_3 = scalar.constant 3 : i8 + %fp4_4 = scalar.constant 4 : i8 + %fp4_5 = scalar.constant 6 : i8 + %fp4_6 = scalar.constant 8 : i8 + %fp4_7 = scalar.constant 12 : i8 + %fp4_8 = scalar.constant 0 : i8 + %fp4_9 = scalar.constant -1 : i8 + %fp4_10 = scalar.constant -2 : i8 + %fp4_11 = scalar.constant -3 : i8 + %fp4_12 = scalar.constant -4 : i8 + %fp4_13 = scalar.constant -6 : i8 + %fp4_14 = scalar.constant -8 : i8 + %fp4_15 = scalar.constant -12 : i8 + %codebook = vector.from_elements %fp4_0, %fp4_1, %fp4_2, %fp4_3, %fp4_4, %fp4_5, %fp4_6, %fp4_7, %fp4_8, %fp4_9, %fp4_10, %fp4_11, %fp4_12, %fp4_13, %fp4_14, %fp4_15 : vector<16xi8> + %lane_block = index.div %lane, %c2 : index + %lane_packet = index.rem %lane, %c2 : index + %packet_offset = index.mul %lane_packet, %c8 : index + %lane_k = index.madd %lane_block, %c32, %packet_offset : index + %lane_sum, %peer_sum = scf.for %base = [%c0 to %k step %c512](%acc = %zero_float : f32, %peer_acc = %zero_float : f32) -> (f32, f32) { + %klo0 = index.add %base, %lane_k : index + %khi0 = index.add %klo0, %c16 : index + %klo = index.assume %klo0 [lt(%klo0, %k), multiple_of(%klo0, 8)] : index + %khi = index.assume %khi0 [lt(%khi0, %k), multiple_of(%khi0, 8)] : index + %block = index.div %klo, %c32 : index + %block_offset = index.scale %block, %block_bytes : index, offset -> offset + %scale_offset = index.add %row_base, %block_offset : offset + %payload_offset = index.add %scale_offset, %one_byte : offset + %scv = buffer.view %weight[%scale_offset] : buffer -> view<1xi8> + %wv = buffer.view %weight[%payload_offset] : buffer -> view<16xi8> + %scale_byte = view.load %scv[%c0] : view<1xi8> -> i8 + %exponent = scalar.extui %scale_byte : i8 to i32 + %normal_bits = scalar.shli %exponent, %scale_shift : i32 + %is_zero = scalar.cmpi eq, %exponent, %zero_i32 : i32 + %is_nan = scalar.cmpi eq, %exponent, %nan_exponent : i32 + %finite_bits = scf.select %is_zero, %scale_subnormal, %normal_bits : i32 + %scale_bits = scf.select %is_nan, %scale_nan, %finite_bits : i32 + %scale = scalar.bitcast %scale_bits : i32 to f32 + %packed_bytes = vector.load %wv[%packet_offset] : view<16xi8> -> vector<8xi8> + %packed = vector.bitcast %packed_bytes : vector<8xi8> to vector<2xi32> + %shifted = vector.shrui %packed, %shift4 : vector<2xi32> + %shifted_bytes = vector.bitcast %shifted : vector<2xi32> to vector<8xi8> + %lo_indices = vector.andi %packed_bytes, %mask4 : vector<8xi8> + %hi_indices = vector.andi %shifted_bytes, %mask4 : vector<8xi8> + %lo = vector.table.lookup %codebook[%lo_indices] : vector<16xi8>, vector<8xi8> -> vector<8xi8> + %hi = vector.table.lookup %codebook[%hi_indices] : vector<16xi8>, vector<8xi8> -> vector<8xi8> + %next = template.apply<@ggml.mxfp4_vector.update>(%klo, %block, %packet_offset, %input_base, %lo, %hi, %scale, %acc, %input) : (index, index, index, offset, vector<8xi8>, vector<8xi8>, f32, f32, buffer) -> (f32) + %peer_next = scf.if %paired -> (f32) { + %value = template.apply<@ggml.mxfp4_vector.update>(%klo, %block, %packet_offset, %peer_base, %lo, %hi, %scale, %peer_acc, %input) : (index, index, index, offset, vector<8xi8>, vector<8xi8>, f32, f32, buffer) -> (f32) + scf.yield %value : f32 + } else { + scf.yield %zero_float : f32 + } + scf.yield %next, %peer_next : f32, f32 + } + template.return %lane_sum, %peer_sum : f32, f32 +} + +kernel.def target(@ggml_mxfp4_dot_wave32) @ggml_mul_mat_id_vector_mxfp4_q8(%token_count: index) { + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %one = index.constant 1 : index + %wave = index.constant 32 : index + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + kernel.launch.config workgroups(%m, %routes, %tokens) workgroup_size(%wave, %one, %one) : index +} launch(%token_count: index, %input: buffer, %ids: buffer, %weight: buffer, %output: buffer) { + %k = config.get @ggml.mxfp4_vector.input_size : index + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %input_routes = config.get @ggml.mxfp4_vector.input_route_count : index + %route_stride = config.get @ggml.mxfp4_vector.route_stride : index + %experts = config.get @ggml.mxfp4_vector.expert_count : index + %row0 = kernel.workgroup.id : index + %route0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %tokens0 = index.assume %token_count [range(%token_count, 1, 5)] : index + %token, %tokens = index.assume %token0, %tokens0 [lt(%token0, %tokens0)] : index, index + %row, %rows = index.assume %row0, %m [lt(%row0, %m)] : index, index + %route, %route_extent = index.assume %route0, %routes [lt(%route0, %routes)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c32 = index.constant 32 : index + %zero_offset = index.constant 0 : offset + %block_bytes = index.constant 17 : offset + %group_bytes = index.constant 144 : offset + %zero_float = scalar.constant 0.0 : f32 + %id_bytes = index.constant 4 : offset + %id_row_bytes = index.scale %route_stride, %id_bytes : index, offset -> offset + %id_offset = index.scale %token, %id_row_bytes : index, offset -> offset + %id_view = buffer.view %ids[%id_offset] : buffer -> view<[%route_extent]xi32> + %expert_i32 = view.load %id_view[%route] : view<[%route_extent]xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %valid = index.cmp ult, %expert0, %experts : index + %sum = scf.if %valid -> (f32) { + %expert, %expert_extent = index.assume %expert0, %experts [lt(%expert0, %experts)] : index, index + %weight_blocks = index.div %k, %c32 : index + %row_bytes = index.scale %weight_blocks, %block_bytes : index, offset -> offset + %flat_row = index.madd %expert, %rows, %row : index + %row_base = index.scale %flat_row, %row_bytes : index, offset -> offset + %input_route = index.rem %route, %input_routes : index + %input_row = index.madd %token, %input_routes, %input_route : index + %groups_per_route = index.div %k, %c128 : index + %route_bytes = index.scale %groups_per_route, %group_bytes : index, offset -> offset + %input_base = index.scale %input_row, %route_bytes : index, offset -> offset + %single = scalar.constant false : i1 + %lane_sum, %unused_peer = template.apply<@ggml.mxfp4_vector.lane_sum>(%single, %row_base, %input_base, %input_base, %input, %weight) : (i1, offset, offset, offset, buffer, buffer) -> (f32, f32) + scf.yield %lane_sum : f32 + } else { + scf.yield %zero_float : f32 + } + %total = kernel.subgroup.reduce %sum : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %y = buffer.view %output[%zero_offset] : buffer -> view<[%tokens]x[%route_extent]x[%rows]xf32> + view.store %total, %y[%token, %route, %row] : f32, view<[%tokens]x[%route_extent]x[%rows]xf32> + } + kernel.return +} + +// Pair equal experts across tokens. Each even occurrence owns the following one. +kernel.def target(@ggml_mxfp4_dot_wave32) @ggml_mxfp4_prepare_pairs(%token_count: index) { + %routes = config.get @ggml.mxfp4_vector.route_count : index + %one = index.constant 1 : index + %group = index.constant 64 : index + %pad = index.constant 63 : index + %slots = index.mul %token_count, %routes : index + %padded = index.add %slots, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%token_count: index, %ids: buffer, %pairs: buffer) { + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %stride = config.get @ggml.mxfp4_vector.route_stride : index + %experts = config.get @ggml.mxfp4_vector.expert_count : index + %slots = index.mul %tokens, %routes : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %zero = index.constant 0 : index + %one = index.constant 1 : index + %two = index.constant 2 : index + %size = index.constant 64 : index + %slot0 = index.madd %group, %size, %lane : index + %valid = index.cmp ult, %slot0, %slots : index + scf.if %valid { + %slot = index.assume %slot0 [range(%slot0, 0, 159), lt(%slot0, %slots)] : index + %token = index.div %slot, %routes : index + %route = index.rem %slot, %routes : index + %four_bytes = index.constant 4 : offset + %row_bytes = index.scale %stride, %four_bytes : index, offset -> offset + %row_base = index.scale %token, %row_bytes : index, offset -> offset + %row = buffer.view %ids[%row_base] : buffer -> view<[%routes]xi32> + %expert = view.load %row[%route] : view<[%routes]xi32> -> i32 + %expert_index = index.cast %expert : i32 to index + %expert_valid = index.cmp ult, %expert_index, %experts : index + %no_peer = scalar.constant -1 : i32 + %owned = scalar.constant -2 : i32 + %partner = scf.if %expert_valid -> (i32) { + %prior_count0, %next0 = scf.for %other = [%zero to %slots step %one](%prior = %zero : index, %next_slot = %slots : index) -> (index, index) { + %other_token = index.div %other, %routes : index + %other_route = index.rem %other, %routes : index + %other_base = index.scale %other_token, %row_bytes : index, offset -> offset + %other_row = buffer.view %ids[%other_base] : buffer -> view<[%routes]xi32> + %other_expert = view.load %other_row[%other_route] : view<[%routes]xi32> -> i32 + %equal = scalar.cmpi eq, %expert, %other_expert : i32 + %before = index.cmp ult, %other, %slot : index + %after = index.cmp ugt, %other, %slot : index + %matching_before = scalar.andi %equal, %before : i1 + %matching_after = scalar.andi %equal, %after : i1 + %increment = scf.select %matching_before, %one, %zero : index + %new_prior = index.add %prior, %increment : index + %candidate = scf.select %matching_after, %other, %next_slot : index + %new_next = index.min %candidate, %next_slot : index + scf.yield %new_prior, %new_next : index, index + } + %prior_count = index.assume %prior_count0 [range(%prior_count0, 0, 159), lt(%prior_count0, %slots)] : index + %next = index.assume %next0 [range(%next0, 0, 160)] : index + %parity = index.rem %prior_count, %two : index + %is_owned = index.cmp ne, %parity, %zero : index + %has_next = index.cmp ult, %next, %slots : index + %next_i32 = index.cast %next : index to i32 + %peer = scf.select %has_next, %next_i32, %no_peer : i32 + %result = scf.select %is_owned, %owned, %peer : i32 + scf.yield %result : i32 + } else { + scf.yield %no_peer : i32 + } + %zero_offset = index.constant 0 : offset + %output = buffer.view %pairs[%zero_offset] : buffer -> view<[%slots]xi32> + view.store %partner, %output[%slot] : i32, view<[%slots]xi32> + } + kernel.return +} + +kernel.def target(@ggml_mxfp4_dot_wave32) @ggml_mul_mat_id_paired_mxfp4_q8(%token_count: index) { + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %one = index.constant 1 : index + %wave = index.constant 32 : index + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + kernel.launch.config workgroups(%m, %routes, %tokens) workgroup_size(%wave, %one, %one) : index +} launch(%token_count: index, %input: buffer, %ids: buffer, %pairs: buffer, %weight: buffer, %output: buffer) { + %k = config.get @ggml.mxfp4_vector.input_size : index + %m = config.get @ggml.mxfp4_vector.output_size : index + %routes = config.get @ggml.mxfp4_vector.route_count : index + %input_routes = config.get @ggml.mxfp4_vector.input_route_count : index + %route_stride = config.get @ggml.mxfp4_vector.route_stride : index + %experts = config.get @ggml.mxfp4_vector.expert_count : index + %row0 = kernel.workgroup.id : index + %route0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %tokens0 = index.assume %token_count [range(%token_count, 1, 5)] : index + %token, %tokens = index.assume %token0, %tokens0 [lt(%token0, %tokens0)] : index, index + %row, %rows = index.assume %row0, %m [lt(%row0, %m)] : index, index + %route, %route_extent = index.assume %route0, %routes [lt(%route0, %routes)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c32 = index.constant 32 : index + %zero_offset = index.constant 0 : offset + %block_bytes = index.constant 17 : offset + %group_bytes = index.constant 144 : offset + %zero_float = scalar.constant 0.0 : f32 + %id_bytes = index.constant 4 : offset + %id_row_bytes = index.scale %route_stride, %id_bytes : index, offset -> offset + %id_offset = index.scale %token, %id_row_bytes : index, offset -> offset + %id_view = buffer.view %ids[%id_offset] : buffer -> view<[%route_extent]xi32> + %expert_i32 = view.load %id_view[%route] : view<[%route_extent]xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %valid = index.cmp ult, %expert0, %experts : index + %slots = index.mul %tokens, %route_extent : index + %slot = index.madd %token, %route_extent, %route : index + %pair_view = buffer.view %pairs[%zero_offset] : buffer -> view<[%slots]xi32> + %peer_i32 = view.load %pair_view[%slot] : view<[%slots]xi32> -> i32 + %owned = scalar.constant -2 : i32 + %active = scalar.cmpi ne, %peer_i32, %owned : i32 + %peer0 = index.cast %peer_i32 : i32 to index + %has_peer = index.cmp ult, %peer0, %slots : index + scf.if %active { + %sum, %peer_sum = scf.if %valid -> (f32, f32) { + %expert, %expert_extent = index.assume %expert0, %experts [lt(%expert0, %experts)] : index, index + %weight_blocks = index.div %k, %c32 : index + %row_bytes = index.scale %weight_blocks, %block_bytes : index, offset -> offset + %flat_row = index.madd %expert, %rows, %row : index + %row_base = index.scale %flat_row, %row_bytes : index, offset -> offset + %input_route = index.rem %route, %input_routes : index + %input_row = index.madd %token, %input_routes, %input_route : index + %groups_per_route = index.div %k, %c128 : index + %route_bytes = index.scale %groups_per_route, %group_bytes : index, offset -> offset + %input_base = index.scale %input_row, %route_bytes : index, offset -> offset + %lane_sum, %lane_peer_sum = scf.if %has_peer -> (f32, f32) { + %peer = index.assume %peer0 [range(%peer0, 0, 159), lt(%peer0, %slots)] : index + %peer_token = index.div %peer, %route_extent : index + %peer_route = index.rem %peer, %route_extent : index + %peer_input_route = index.rem %peer_route, %input_routes : index + %peer_input_row = index.madd %peer_token, %input_routes, %peer_input_route : index + %peer_base = index.scale %peer_input_row, %route_bytes : index, offset -> offset + %paired = scalar.constant true : i1 + %owner_sum, %partner_sum = template.apply<@ggml.mxfp4_vector.lane_sum>(%paired, %row_base, %input_base, %peer_base, %input, %weight) : (i1, offset, offset, offset, buffer, buffer) -> (f32, f32) + scf.yield %owner_sum, %partner_sum : f32, f32 + } else { + %single = scalar.constant false : i1 + %single_sum, %unused_peer_sum = template.apply<@ggml.mxfp4_vector.lane_sum>(%single, %row_base, %input_base, %input_base, %input, %weight) : (i1, offset, offset, offset, buffer, buffer) -> (f32, f32) + scf.yield %single_sum, %unused_peer_sum : f32, f32 + } + scf.yield %lane_sum, %lane_peer_sum : f32, f32 + } else { + scf.yield %zero_float, %zero_float : f32, f32 + } + %total = kernel.subgroup.reduce %sum : f32 + %peer_total = kernel.subgroup.reduce %peer_sum : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %y = buffer.view %output[%zero_offset] : buffer -> view<[%tokens]x[%route_extent]x[%rows]xf32> + view.store %total, %y[%token, %route, %row] : f32, view<[%tokens]x[%route_extent]x[%rows]xf32> + scf.if %has_peer { + %peer = index.assume %peer0 [range(%peer0, 0, 159), lt(%peer0, %slots)] : index + %peer_token = index.div %peer, %route_extent : index + %peer_route = index.rem %peer, %route_extent : index + view.store %peer_total, %y[%peer_token, %peer_route, %row] : f32, view<[%tokens]x[%route_extent]x[%rows]xf32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argmax_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argmax_f32.loom new file mode 100644 index 000000000000..d8e55c55950b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argmax_f32.loom @@ -0,0 +1,39 @@ +config.decl @ggml.argmax.width : %value: index where [range(%value, 1, 1048576)] + +// GGML chooses the last index when finite maxima tie. +kernel.def @ggml_argmax_f32(%row_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %bounded_rows = index.assume %row_count [range(%row_count, 1, 2048)] : index + kernel.launch.config workgroups(%bounded_rows, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %count = index.assume %row_count [range(%row_count, 1, 2048)] : index + %width = config.get @ggml.argmax.width : index + %row0 = kernel.workgroup.id : index + %row, %row_extent = index.assume %row0, %count [range(%row0, 0, 2047), lt(%row0, %count)] : index, index + %lane = kernel.workitem.id : index + %step = index.constant 256 : index + %c0 = index.constant 0 : index + %zero = index.constant 0 : offset + %zero_i32 = scalar.constant 0 : i32 + %negative_inf = scalar.constant -inf : f32 + %source = buffer.view %input[%zero] : buffer -> view<[%count]x[%width]xf32> + %dest = buffer.view %output[%zero] : buffer -> view<[%count]xi32> + %partial, %partial_index = scf.for %column = [%lane to %width step %step](%best = %negative_inf : f32, %best_index = %zero_i32 : i32) -> (f32, i32) { + %value = view.load %source[%row, %column] : view<[%count]x[%width]xf32> -> f32 + %take = scalar.cmpf oge, %value, %best : f32 + %column_i32 = index.cast %column : index to i32 + %next = scf.select %take, %value, %best : f32 + %next_index = scf.select %take, %column_i32, %best_index : i32 + scf.yield %next, %next_index : f32, i32 + } + %maximum = kernel.workgroup.reduce %partial : f32 + %matches = scalar.cmpf oeq, %partial, %maximum : f32 + %candidate = scf.select %matches, %partial_index, %zero_i32 : i32 + %winner = kernel.workgroup.reduce %candidate : i32 + %writes = index.cmp eq, %lane, %c0 : index + scf.if %writes { + view.store %winner, %dest[%row] : i32, view<[%count]xi32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argsort_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argsort_f32.loom new file mode 100644 index 000000000000..591dee6461cb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/argsort_f32.loom @@ -0,0 +1,55 @@ +config.decl @ggml.argsort.width : %value: index where [range(%value, 0, 2048)] +config.decl @ggml.argsort.output_width : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.argsort.descending : %value: index where [range(%value, 0, 2048)] +kernel.def @ggml_argsort_f32(%row_count: index) { + %width = config.get @ggml.argsort.width : index + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %width, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %row_count, %one) workgroup_size(%group, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %width0 = config.get @ggml.argsort.width : index + %width = index.assume %width0 [range(%width0, 1, 2048)] : index + %count = index.assume %row_count [range(%row_count, 1, 1048576)] : index + %desc = config.get @ggml.argsort.descending : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %row0 = kernel.workgroup.id : index + %row, %extent = index.assume %row0, %count [range(%row0, 0, 1048575), lt(%row0, %count)] : index, index + %group_size = index.constant 256 : index + %column0 = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %column0, %width : index + scf.if %valid { + %column, %width_extent = index.assume %column0, %width [range(%column0, 0, 2047), lt(%column0, %width)] : index, index + %zero = index.constant 0 : index + %one = index.constant 1 : index + %offset = index.constant 0 : offset + %descending = index.cmp ne, %desc, %zero : index + %output_width = config.get @ggml.argsort.output_width : index + %x = buffer.view %input[%offset] : buffer -> view<[%count]x[%width]xf32> + %y = buffer.view %output[%offset] : buffer -> view<[%count]x[%output_width]xi32> + %value = view.load %x[%row, %column] : view<[%count]x[%width]xf32> -> f32 + %rank0 = scf.for %j = [%zero to %width step %one](%rank = %zero : index) -> (index) { + %other = view.load %x[%row, %j] : view<[%count]x[%width]xf32> -> f32 + %greater = scalar.cmpf ogt, %other, %value : f32 + %less = scalar.cmpf olt, %other, %value : f32 + %equal = scalar.cmpf oeq, %other, %value : f32 + %before = index.cmp ult, %j, %column : index + %tie = scalar.andi %equal, %before : i1 + %ordered = scf.select %descending, %greater, %less : i1 + %precedes = scalar.ori %ordered, %tie : i1 + %increment = scf.select %precedes, %one, %zero : index + %next = index.add %rank, %increment : index + scf.yield %next : index + } + %rank, %rank_extent = index.assume %rank0, %width [range(%rank0, 0, 2047), lt(%rank0, %width)] : index, index + %index = index.cast %column : index to i32 + %selected = index.cmp ult, %rank, %output_width : index + scf.if %selected { + view.store %index, %y[%row, %rank] : i32, view<[%count]x[%output_width]xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_32.loom new file mode 100644 index 000000000000..10bcaa1aa5f0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_32.loom @@ -0,0 +1,64 @@ +// Bit-preserving batched row gather for F32 and I32 tensors. +config.decl @ggml.get_rows_32.width : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.rows : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.selected : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.batch1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.source_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.ids_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.source_stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.source_stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.source_stride3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.ids_stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.get_rows_32.ids_stride2 : %value: index where [range(%value, 1, 1073741824)] + +kernel.def @ggml_get_rows_32(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %groups = index.div %rounded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %source: buffer, %ids: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %width = config.get @ggml.get_rows_32.width : index + %rows = config.get @ggml.get_rows_32.rows : index + %selected = config.get @ggml.get_rows_32.selected : index + %batch1 = config.get @ggml.get_rows_32.batch1 : index + %source_span = config.get @ggml.get_rows_32.source_span : index + %ids_span = config.get @ggml.get_rows_32.ids_span : index + %source_stride1 = config.get @ggml.get_rows_32.source_stride1 : index + %source_stride2 = config.get @ggml.get_rows_32.source_stride2 : index + %source_stride3 = config.get @ggml.get_rows_32.source_stride3 : index + %ids_stride1 = config.get @ggml.get_rows_32.ids_stride1 : index + %ids_stride2 = config.get @ggml.get_rows_32.ids_stride2 : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %size = index.constant 256 : index + %linear = index.madd %group, %size, %lane : index + %valid = index.cmp ult, %linear, %count : index + %zero = index.constant 0 : offset + %source_view = buffer.view %source[%zero] : buffer -> view<[%source_span]xi32> + %ids_view = buffer.view %ids[%zero] : buffer -> view<[%ids_span]xi32> + %output_view = buffer.view %output[%zero] : buffer -> view<[%count]xi32> + scf.if %valid { + %column = index.rem %linear, %width : index + %row_flat = index.div %linear, %width : index + %selected_row = index.rem %row_flat, %selected : index + %batch_flat = index.div %row_flat, %selected : index + %batch_a = index.rem %batch_flat, %batch1 : index + %batch_b = index.div %batch_flat, %batch1 : index + %ids_index_a = index.madd %batch_a, %ids_stride1, %selected_row : index + %ids_index0 = index.madd %batch_b, %ids_stride2, %ids_index_a : index + %ids_index, %ids_extent = index.assume %ids_index0, %ids_span [range(%ids_index0, 0, 1073741823), lt(%ids_index0, %ids_span)] : index, index + %row_i32 = view.load %ids_view[%ids_index] : view<[%ids_span]xi32> -> i32 + %row0 = index.cast %row_i32 : i32 to index + %row, %rows_extent = index.assume %row0, %rows [range(%row0, 0, 1073741823), lt(%row0, %rows)] : index, index + %source_row = index.madd %row, %source_stride1, %column : index + %source_batch_a = index.madd %batch_a, %source_stride2, %source_row : index + %source_index0 = index.madd %batch_b, %source_stride3, %source_batch_a : index + %source_index, %source_extent = index.assume %source_index0, %source_span [range(%source_index0, 0, 1073741823), lt(%source_index0, %source_span)] : index, index + %value = view.load %source_view[%source_index] : view<[%source_span]xi32> -> i32 + view.store %value, %output_view[%linear] : i32, view<[%count]xi32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom index 205339245cb8..3ae48a1537db 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom @@ -1,3 +1,6 @@ +config.decl @ggml.set_rows.index_format : %value: index where [range(%value, 32, 64)] +template.decl @ggml.rows.cache_publish.safe_i32_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) + // Generic GGML SET_ROWS for contiguous 2D row caches. template.decl @ggml.set_rows.body(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) @@ -73,7 +76,16 @@ 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 - %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) + %index_format = config.get @ggml.set_rows.index_format : index + %index32 = index.constant 32 : index + %uses_i32 = index.cmp eq, %index_format, %index32 : index + %valid_row, %cache_row = scf.if %uses_i32 -> (i1, index) { + %valid, %row = template.apply<@ggml.rows.cache_publish.safe_i32_row>(%token, %bounded_token_count, %bounded_cache_row_count, %indices_noalias) : (index, index, index, buffer) -> (i1, index) + scf.yield %valid, %row : i1, index + } else { + %valid, %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) + scf.yield %valid, %row : 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.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) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows_scalar.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows_scalar.loom new file mode 100644 index 000000000000..5fdedec84d4d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows_scalar.loom @@ -0,0 +1,109 @@ +config.decl @ggml.set_rows_scalar.nr : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.in1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.in2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.is0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.is1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.is2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.rs1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.rs2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.rs3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.cs1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.cs2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.cs3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.index_bits : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.element_bits : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.index_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.source_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.cache_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.set_rows_scalar.cache_rows : %value: index where [range(%value, 1, 1073741824)] +kernel.def @ggml_set_rows_scalar(%element_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + %pad = index.constant 255 : index + %padded = index.add %element_count, %pad : index + %groups = index.div %padded, %group : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%element_count: index, %rows: buffer, %indices: buffer, %cache: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %nr = config.get @ggml.set_rows_scalar.nr : index + %ne2 = config.get @ggml.set_rows_scalar.ne2 : index + %in1 = config.get @ggml.set_rows_scalar.in1 : index + %in2 = config.get @ggml.set_rows_scalar.in2 : index + %is0 = config.get @ggml.set_rows_scalar.is0 : index + %is1 = config.get @ggml.set_rows_scalar.is1 : index + %is2 = config.get @ggml.set_rows_scalar.is2 : index + %rs1 = config.get @ggml.set_rows_scalar.rs1 : index + %rs2 = config.get @ggml.set_rows_scalar.rs2 : index + %rs3 = config.get @ggml.set_rows_scalar.rs3 : index + %cs1 = config.get @ggml.set_rows_scalar.cs1 : index + %cs2 = config.get @ggml.set_rows_scalar.cs2 : index + %cs3 = config.get @ggml.set_rows_scalar.cs3 : index + %index_bits = config.get @ggml.set_rows_scalar.index_bits : index + %element_bits = config.get @ggml.set_rows_scalar.element_bits : index + %index_span = config.get @ggml.set_rows_scalar.index_span : index + %source_span = config.get @ggml.set_rows_scalar.source_span : index + %cache_span = config.get @ggml.set_rows_scalar.cache_span : index + %cache_rows = config.get @ggml.set_rows_scalar.cache_rows : index + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %group_size = index.constant 256 : index + %linear = index.madd %group, %group_size, %lane : index + %valid = index.cmp ult, %linear, %count : index + scf.if %valid { + %row = index.rem %linear, %nr : index + %rest = index.div %linear, %nr : index + %b2 = index.rem %rest, %ne2 : index + %b3 = index.div %rest, %ne2 : index + %ib1 = index.rem %b2, %in1 : index + %ib2 = index.rem %b3, %in2 : index + %ip0 = index.mul %row, %is0 : index + %ip1 = index.madd %ib1, %is1, %ip0 : index + %ip2 = index.madd %ib2, %is2, %ip1 : index + %ip, %ie = index.assume %ip2, %index_span [range(%ip2, 0, 1073741823), lt(%ip2, %index_span)] : index, index + %zero = index.constant 0 : offset + %bits32 = index.constant 32 : index + %is32 = index.cmp eq, %index_bits, %bits32 : index + %raw = scf.if %is32 -> (i64) { + %iv = buffer.view %indices[%zero] : buffer -> view<[%index_span]xi32> + %v = view.load %iv[%ip] : view<[%index_span]xi32> -> i32 + %wide = scalar.extsi %v : i32 to i64 + scf.yield %wide : i64 + } else { + %iv = buffer.view %indices[%zero] : buffer -> view<[%index_span]xi64> + %v = view.load %iv[%ip] : view<[%index_span]xi64> -> i64 + scf.yield %v : i64 + } + %raw_zero = scalar.constant 0 : i64 + %raw_limit = index.cast %cache_rows : index to i64 + %nonnegative = scalar.cmpi sge, %raw, %raw_zero : i64 + %below = scalar.cmpi slt, %raw, %raw_limit : i64 + %valid_row = scalar.andi %nonnegative, %below : i1 + scf.if %valid_row { + %cache_row0 = index.cast %raw : i64 to index + %cache_row, %cre = index.assume %cache_row0, %cache_rows [range(%cache_row0, 0, 1073741823), lt(%cache_row0, %cache_rows)] : index, index + %rp0 = index.mul %row, %rs1 : index + %rp1 = index.madd %b2, %rs2, %rp0 : index + %rp2 = index.madd %b3, %rs3, %rp1 : index + %rp, %re = index.assume %rp2, %source_span [range(%rp2, 0, 1073741823), lt(%rp2, %source_span)] : index, index + %cp0 = index.mul %cache_row, %cs1 : index + %cp1 = index.madd %b2, %cs2, %cp0 : index + %cp2 = index.madd %b3, %cs3, %cp1 : index + %cp, %ce = index.assume %cp2, %cache_span [range(%cp2, 0, 1073741823), lt(%cp2, %cache_span)] : index, index + %bits16 = index.constant 16 : index + %is16 = index.cmp eq, %element_bits, %bits16 : index + scf.if %is16 { + %src = buffer.view %rows[%zero] : buffer -> view<[%source_span]xi16> + %dst = buffer.view %cache[%zero] : buffer -> view<[%cache_span]xi16> + %value = view.load %src[%rp] : view<[%source_span]xi16> -> i16 + view.store %value, %dst[%cp] : i16, view<[%cache_span]xi16> + } else { + %src = buffer.view %rows[%zero] : buffer -> view<[%source_span]xi32> + %dst = buffer.view %cache[%zero] : buffer -> view<[%cache_span]xi32> + %value = view.load %src[%rp] : view<[%source_span]xi32> -> i32 + view.store %value, %dst[%cp] : i32, view<[%cache_span]xi32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/softmax_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/softmax_f32.loom new file mode 100644 index 000000000000..33de27d0daf5 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/softmax_f32.loom @@ -0,0 +1,46 @@ +config.decl @ggml.softmax.width : %value: index where [range(%value, 1, 32768)] +config.decl @ggml.softmax.scale : f32 + +kernel.def @ggml_softmax_f32(%row_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + kernel.launch.config workgroups(%row_count, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %rows = index.assume %row_count [range(%row_count, 1, 134217728)] : index + %width = config.get @ggml.softmax.width : index + %scale = config.get @ggml.softmax.scale : f32 + %row0 = kernel.workgroup.id : index + %row, %extent = index.assume %row0, %rows [range(%row0, 0, 134217727), lt(%row0, %rows)] : index, index + %lane = kernel.workitem.id : index + %step = index.constant 256 : index + %zero = index.constant 0 : offset + %zero_f32 = scalar.constant 0.0 : f32 + %neg_inf = scalar.constant -inf : f32 + %source = buffer.view %input[%zero] : buffer -> view<[%rows]x[%width]xf32> + %dest = buffer.view %output[%zero] : buffer -> view<[%rows]x[%width]xf32> + %partial_max = scf.for %column = [%lane to %width step %step](%acc = %neg_inf : f32) -> (f32) { + %value = view.load %source[%row, %column] : view<[%rows]x[%width]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %next = scalar.maxnumf %acc, %scaled : f32 + scf.yield %next : f32 + } + %maximum = kernel.workgroup.reduce %partial_max : f32 + %partial_sum = scf.for %column = [%lane to %width step %step](%acc = %zero_f32 : f32) -> (f32) { + %value = view.load %source[%row, %column] : view<[%rows]x[%width]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %shifted = scalar.subf %scaled, %maximum : f32 + %exp = scalar.expf %shifted : f32 + %next = scalar.addf %acc, %exp : f32 + scf.yield %next : f32 + } + %sum = kernel.workgroup.reduce %partial_sum : f32 + scf.for %column = [%lane to %width step %step] { + %value = view.load %source[%row, %column] : view<[%rows]x[%width]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %shifted = scalar.subf %scaled, %maximum : f32 + %exp = scalar.expf %shifted : f32 + %result = scalar.divf %exp, %sum : f32 + view.store %result, %dest[%row, %column] : f32, view<[%rows]x[%width]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/sum_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/sum_rows_f32.loom new file mode 100644 index 000000000000..51bf3b8e556f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/sum_rows_f32.loom @@ -0,0 +1,53 @@ +config.decl @ggml.sum_rows.width : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.ne1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.source_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.sum_rows.stride3 : %value: index where [range(%value, 1, 1073741824)] + +kernel.def @ggml_sum_rows_f32(%row_count: index) { + %one = index.constant 1 : index + %group = index.constant 256 : index + kernel.launch.config workgroups(%row_count, %one, %one) workgroup_size(%group, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %count = index.assume %row_count [range(%row_count, 1, 1048576)] : index + %width = config.get @ggml.sum_rows.width : index + %ne1 = config.get @ggml.sum_rows.ne1 : index + %ne2 = config.get @ggml.sum_rows.ne2 : index + %source_span = config.get @ggml.sum_rows.source_span : index + %stride0 = config.get @ggml.sum_rows.stride0 : index + %stride1 = config.get @ggml.sum_rows.stride1 : index + %stride2 = config.get @ggml.sum_rows.stride2 : index + %stride3 = config.get @ggml.sum_rows.stride3 : index + %row0 = kernel.workgroup.id : index + %row, %row_extent = index.assume %row0, %count [range(%row0, 0, 1048575), lt(%row0, %count)] : index, index + %lane = kernel.workitem.id : index + %step = index.constant 256 : index + %c0 = index.constant 0 : index + %zero = index.constant 0 : offset + %zero_f32 = scalar.constant 0.0 : f32 + %i1 = index.rem %row, %ne1 : index + %rest = index.div %row, %ne1 : index + %i2 = index.rem %rest, %ne2 : index + %i3 = index.div %rest, %ne2 : index + %offset1 = index.mul %i1, %stride1 : index + %offset2 = index.madd %i2, %stride2, %offset1 : index + %base = index.madd %i3, %stride3, %offset2 : index + %source = buffer.view %input[%zero] : buffer -> view<[%source_span]xf32> + %dest = buffer.view %output[%zero] : buffer -> view<[%count]xf32> + %partial = scf.for %column = [%lane to %width step %step](%acc = %zero_f32 : f32) -> (f32) { + %offset0 = index.madd %column, %stride0, %base : index + %offset, %extent = index.assume %offset0, %source_span [range(%offset0, 0, 1073741823), lt(%offset0, %source_span)] : index, index + %value = view.load %source[%offset] : view<[%source_span]xf32> -> f32 + %next = scalar.addf %acc, %value : f32 + scf.yield %next : f32 + } + %sum = kernel.workgroup.reduce %partial : f32 + %writes = index.cmp eq, %lane, %c0 : index + scf.if %writes { + view.store %sum, %dest[%row] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/top_k_radix_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/top_k_radix_f32.loom new file mode 100644 index 000000000000..3c8e902d412f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/top_k_radix_f32.loom @@ -0,0 +1,231 @@ +amdgpu.target @ggml_top_k_radix_wave64 {subgroup_size = 64} + +config.decl @ggml.top_k.width : %value: index where [range(%value, 1, 131072)] + +config.decl @ggml.top_k.count : %value: index where [range(%value, 1, 512)] + +func.def inline @ggml_top_k_float_key(%value: f32) -> (i32) { + %zero = scalar.constant 0 : i32 + %ones = scalar.constant -1 : i32 + %sign = scalar.constant -2147483648 : i32 + %zero_float = scalar.constant 0.0 : f32 + %bits = scalar.bitcast %value : f32 to i32 + %negative = scalar.cmpi slt, %bits, %zero : i32 + %flip = scf.select %negative, %ones, %sign : i32 + %ordered = scalar.xori %bits, %flip : i32 + %is_zero = scalar.cmpf oeq, %value, %zero_float : f32 + %is_nan = scalar.cmpf uno, %value, %value : f32 + %canonical = scf.select %is_zero, %sign, %ordered : i32 + %key = scf.select %is_nan, %zero, %canonical : i32 + func.return %key : i32 +} + +config.def @ggml.top_k.stage_width = 1 : index + +config.def @ggml.top_k.first_pass = 1 : index + +config.def @ggml.top_k.last_pass = 1 : index + +template.decl @ggml.top_k.radix(%partitioned: i1, %row_count: index, %input: buffer, %input_ids: buffer, %output: buffer) + +template.def<@ggml.top_k.radix> device @ggml_top_k_radix(%partitioned: i1, %row_count: index, %input: buffer, %input_ids: buffer, %output: buffer) { + %rows = index.assume %row_count [range(%row_count, 1, 2048)] : index + %source_width = config.get @ggml.top_k.width : index + %requested_count = config.get @ggml.top_k.count : index + %row0 = kernel.workgroup.id : index + %row = index.assume %row0 [lt(%row0, %rows)] : index + %thread = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c24 = index.constant 24 : index + %c256 = index.constant 256 : index + %width, %input_width, %output_width, %partition_base, %output_base, %first_pass, %last_pass = scf.if %partitioned -> (index, index, index, index, index, i1, i1) { + %stage_width = config.get @ggml.top_k.stage_width : index + %first = config.get @ggml.top_k.first_pass : index + %last = config.get @ggml.top_k.last_pass : index + %is_first = index.cmp ne, %first, %c0 : index + %is_last = index.cmp ne, %last, %c0 : index + %padding = index.constant 255 : index + %padded_width = index.add %stage_width, %padding : index + %partition_count = index.div %padded_width, %c256 : index + %partition0 = kernel.workgroup.id : index + %partition = index.assume %partition0 [lt(%partition0, %partition_count)] : index + %base = index.mul %partition, %c256 : index + %remaining = index.sub %stage_width, %base : index + %extent = index.min %remaining, %c256 : index + %full_partitions = index.div %stage_width, %c256 : index + %tail = index.rem %stage_width, %c256 : index + %tail_count = index.min %tail, %requested_count : index + %result_width = index.madd %full_partitions, %requested_count, %tail_count : index + %result_base = index.mul %partition, %requested_count : index + scf.yield %extent, %stage_width, %result_width, %base, %result_base, %is_first, %is_last : index, index, index, index, index, i1, i1 + } else { + %yes = scalar.constant true : i1 + scf.yield %source_width, %source_width, %requested_count, %c0, %c0, %yes, %yes : index, index, index, index, index, i1, i1 + } + %count = index.min %requested_count, %width : index + %offset = index.constant 0 : offset + %hist_bytes = index.constant 1024 : offset + %state_bytes = index.constant 8 : offset + %hist_buffer = buffer.alloca align(16) %hist_bytes : buffer + %state_buffer = buffer.alloca align(8) %state_bytes : buffer + %hist = buffer.view %hist_buffer[%offset] : buffer -> view<256xi32> + %state = buffer.view %state_buffer[%offset] : buffer -> view<2xi32> + %source = buffer.view %input[%offset] : buffer -> view<[%rows]x[%source_width]xf32> + %target = buffer.view %output[%offset] : buffer -> view<[%rows]x[%output_width]xi32> + %indices = buffer.view %input_ids[%offset] : buffer -> view<[%rows]x[%input_width]xi32> + %zero_i32 = scalar.constant 0 : i32 + %one_i32 = scalar.constant 1 : i32 + %ones_i32 = scalar.constant -1 : i32 + %byte_mask = scalar.constant 255 : i32 + %count_i32 = index.cast %count : index to i32 + %threshold, %ties_needed = scf.for %pass = [%c0 to %c4 step %c1](%prefix = %zero_i32 : i32, %desired = %count_i32 : i32) -> (i32, i32) unroll { + view.store %zero_i32, %hist[%thread] : i32, view<256xi32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %bit = index.mul %pass, %c8 : index + %shift_index = index.sub %c24, %bit : index + %shift = index.cast %shift_index : index to i32 + %first = index.cmp eq, %pass, %c0 : index + %high_mask = scf.if %first -> (i32) { + scf.yield %zero_i32 : i32 + } else { + %high_shift_index = index.add %shift_index, %c8 : index + %high_shift = index.cast %high_shift_index : index to i32 + %mask = scalar.shli %ones_i32, %high_shift : i32 + scf.yield %mask : i32 + } + scf.for %column = [%thread to %width step %c256] { + %input_column0 = index.add %partition_base, %column : index + %input_column = index.assume %input_column0 [lt(%input_column0, %input_width)] : index + %source_column0 = scf.if %first_pass -> (index) { + scf.yield %input_column : index + } else { + %id = view.load %indices[%row, %input_column] : view<[%rows]x[%input_width]xi32> -> i32 + %column = index.cast %id : i32 to index + scf.yield %column : index + } + %source_column = index.assume %source_column0 [range(%source_column0, 0, 131071), lt(%source_column0, %source_width)] : index + %value = view.load %source[%row, %source_column] : view<[%rows]x[%source_width]xf32> -> f32 + %key = func.call @ggml_top_k_float_key(%value) : (f32) -> (i32) + %high = scalar.andi %key, %high_mask : i32 + %candidate = scalar.cmpi eq, %high, %prefix : i32 + scf.if %candidate { + %shifted = scalar.shrui %key, %shift : i32 + %bucket_i32 = scalar.andi %shifted, %byte_mask : i32 + %bucket0 = index.cast %bucket_i32 : i32 to index + %bucket = index.assume %bucket0 [range(%bucket0, 0, 255)] : index + view.atomic.reduce %one_i32, %hist[%bucket] {ordering = relaxed, scope = workgroup} : i32, view<256xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %bucket_count = view.load %hist[%thread] : view<256xi32> -> i32 + %above = kernel.workgroup.scan %bucket_count {direction = reverse, mode = exclusive} : i32 + %end = scalar.addi %above, %bucket_count : i32 + %starts_before = scalar.cmpi ult, %above, %desired : i32 + %ends_after = scalar.cmpi uge, %end, %desired : i32 + %selected = scalar.andi %starts_before, %ends_after : i1 + scf.if %selected { + %bucket_i32 = index.cast %thread : index to i32 + %bucket_bits = scalar.shli %bucket_i32, %shift : i32 + %next_prefix = scalar.ori %prefix, %bucket_bits : i32 + %next_desired = scalar.subi %desired, %above : i32 + view.store %next_prefix, %state[0] : i32, view<2xi32> + view.store %next_desired, %state[1] : i32, view<2xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_prefix = view.load %state[0] : view<2xi32> -> i32 + %next_desired = view.load %state[1] : view<2xi32> -> i32 + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_prefix, %next_desired : i32, i32 + } + %above_total = scalar.subi %count_i32, %ties_needed : i32 + %counts_zero = vector.constant 0 : vector<2xi32> + %final_counts = scf.for %base = [%c0 to %width step %c256](%counts = %counts_zero : vector<2xi32>) -> (vector<2xi32>) { + %position = index.add %base, %thread : index + %valid = index.cmp ult, %position, %width : index + %key = scf.if %valid -> (i32) { + %column = index.assume %position [lt(%position, %width)] : index + %input_column0 = index.add %partition_base, %column : index + %input_column = index.assume %input_column0 [lt(%input_column0, %input_width)] : index + %source_column0 = scf.if %first_pass -> (index) { + scf.yield %input_column : index + } else { + %id = view.load %indices[%row, %input_column] : view<[%rows]x[%input_width]xi32> -> i32 + %column = index.cast %id : i32 to index + scf.yield %column : index + } + %source_column = index.assume %source_column0 [range(%source_column0, 0, 131071), lt(%source_column0, %source_width)] : index + %value = view.load %source[%row, %source_column] : view<[%rows]x[%source_width]xf32> -> f32 + %key = func.call @ggml_top_k_float_key(%value) : (f32) -> (i32) + scf.yield %key : i32 + } else { + scf.yield %zero_i32 : i32 + } + %greater0 = scalar.cmpi ugt, %key, %threshold : i32 + %equal0 = scalar.cmpi eq, %key, %threshold : i32 + %greater = scalar.andi %valid, %greater0 : i1 + %equal = scalar.andi %valid, %equal0 : i1 + %greater_count = scf.select %greater, %one_i32, %zero_i32 : i32 + %equal_count = scf.select %equal, %one_i32, %zero_i32 : i32 + %increments = vector.from_elements %greater_count, %equal_count : vector<2xi32> + %prefix = kernel.workgroup.scan %increments {direction = forward, mode = exclusive} : vector<2xi32> + %block_counts = kernel.workgroup.reduce %increments : vector<2xi32> + %ranks = vector.addi %counts, %prefix : vector<2xi32> + %greater_rank = vector.extract %ranks[0] : vector<2xi32> -> i32 + %equal_rank = vector.extract %ranks[1] : vector<2xi32> -> i32 + %tie_needed = scalar.cmpi ult, %equal_rank, %ties_needed : i32 + %take_tie = scalar.andi %equal, %tie_needed : i1 + %selected = scalar.ori %greater, %take_tie : i1 + %tie_slot = scalar.addi %above_total, %equal_rank : i32 + %final_slot = scf.select %greater, %greater_rank, %tie_slot : i32 + // Keep intermediate IDs in source order so later partitions preserve tie ordering. + %prior_ties = scalar.minui %equal_rank, %ties_needed : i32 + %stable_slot = scalar.addi %greater_rank, %prior_ties : i32 + %slot_i32 = scf.select %last_pass, %final_slot, %stable_slot : i32 + scf.if %selected { + %slot0 = index.cast %slot_i32 : i32 to index + %slot = index.assume %slot0 [range(%slot0, 0, 511), lt(%slot0, %count)] : index + %input_column0 = index.add %partition_base, %position : index + %input_column = index.assume %input_column0 [lt(%input_column0, %input_width)] : index + %id = scf.if %first_pass -> (i32) { + %id = index.cast %input_column : index to i32 + scf.yield %id : i32 + } else { + %id = view.load %indices[%row, %input_column] : view<[%rows]x[%input_width]xi32> -> i32 + scf.yield %id : i32 + } + %result_column0 = index.add %output_base, %slot : index + %result_column = index.assume %result_column0 [lt(%result_column0, %output_width)] : index + view.store %id, %target[%row, %result_column] : i32, view<[%rows]x[%output_width]xi32> + } + %next_counts = vector.addi %counts, %block_counts : vector<2xi32> + scf.yield %next_counts : vector<2xi32> + } + template.return +} + +kernel.def target(@ggml_top_k_radix_wave64) @ggml_top_k_radix_f32(%row_count: index) { + %one = index.constant 1 : index + %threads = index.constant 256 : index + kernel.launch.config workgroups(%one, %row_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%row_count: index, %input: buffer, %output: buffer) { + %partitioned = scalar.constant false : i1 + template.apply<@ggml.top_k.radix>(%partitioned, %row_count, %input, %output, %output) : (i1, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_top_k_radix_wave64) @ggml_top_k_radix_partitions_f32(%row_count: index) { + %one = index.constant 1 : index + %threads = index.constant 256 : index + %padding = index.constant 255 : index + %width = config.get @ggml.top_k.stage_width : index + %padded = index.add %width, %padding : index + %partitions = index.div %padded, %threads : index + kernel.launch.config workgroups(%partitions, %row_count, %one) workgroup_size(%threads, %one, %one) : index +} launch(%row_count: index, %input: buffer, %input_ids: buffer, %output: buffer) { + %partitioned = scalar.constant true : i1 + template.apply<@ggml.top_k.radix>(%partitioned, %row_count, %input, %input_ids, %output) : (i1, index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index d62b2be3ccfb..4a0b51cfec46 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -684,18 +684,25 @@ GraphProgramLookup GraphProgramCache::build_from_imported(const ggml_cgraph & g } if (graph.uid == 0) { - result.uncached_program = std::move(program); - result.program = result.uncached_program.get(); + result.program = std::move(program); result.match = std::move(match); return result; } - GraphProgram * cached_program = program.get(); + std::shared_ptr cached_program = std::move(program); { std::lock_guard lock(mutex_); validated_matches_.clear(); - programs_[graph.uid] = std::move(program); - cached_program = programs_[graph.uid].get(); + auto retired = programs_.find(graph.uid); + if (retired == programs_.end() && programs_.size() >= 32) { + retired = std::min_element(programs_.begin(), programs_.end(), [](const auto & a, const auto & b) { + return a.second.last_used < b.second.last_used; + }); + } + if (retired != programs_.end()) { + programs_.erase(retired); + } + programs_[graph.uid] = { cached_program, ++access_tick_ }; last_program_ = cached_program; ++stats_.builds; } @@ -712,7 +719,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, const bool disable_fast_path = environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_UID_FAST_PATH"); const bool validate_fast_path = environment_flag_enabled("GGML_HRX_VALIDATE_GRAPH_UID_CACHE"); const bool allow_structural_reuse = std::getenv("GGML_HRX_DUMP_COMMAND_PROGRAM_DIR") == nullptr; - GraphProgram * cached_program = nullptr; + std::shared_ptr cached_program; { std::lock_guard lock(mutex_); if (!disable_fast_path && last_program_ != nullptr && last_program_->uid() == graph.uid && @@ -720,8 +727,8 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, cached_program = last_program_; } else { const auto found = programs_.find(graph.uid); - if (found != programs_.end() && found->second->target() == target) { - cached_program = found->second.get(); + if (found != programs_.end() && found->second.program->target() == target) { + cached_program = found->second.program; last_program_ = cached_program; } if (cached_program == nullptr && !disable_fast_path && allow_structural_reuse) { @@ -732,6 +739,9 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, } } } + if (cached_program != nullptr) { + programs_.at(cached_program->uid()).last_used = ++access_tick_; + } } if (cached_program != nullptr) { const bool bind_external = @@ -766,8 +776,8 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, // Physical bindings have a separate prepared-program cache key. if (allow_structural_reuse) { std::lock_guard lock(mutex_); - for (const auto & entry : programs_) { - GraphProgram * candidate = entry.second.get(); + for (auto & entry : programs_) { + const auto & candidate = entry.second.program; if (candidate == cached_program || candidate->target() != target) { continue; } @@ -776,6 +786,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, continue; } last_program_ = candidate; + entry.second.last_used = ++access_tick_; if (!disable_fast_path) { if (validated_matches_.size() >= 128) { validated_matches_.clear(); diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index 54c41ca6873c..fa1ef30528d6 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -113,8 +113,7 @@ struct GraphProgramCacheStats { }; struct GraphProgramLookup { - GraphProgram * program = nullptr; - std::unique_ptr uncached_program; + std::shared_ptr program; GraphProgramMatch match; Status status; @@ -144,10 +143,15 @@ class GraphProgramCache { private: struct ValidatedGraphMatch { - GraphProgram * program = nullptr; + std::shared_ptr program; const ggml_tensor * const * nodes = nullptr; }; + struct Entry { + std::shared_ptr program; + uint64_t last_used = 0; + }; + GraphProgramLookup build_from_imported(const ggml_cgraph & graph, Graph && imported_graph, const KernelCorpus & corpus, @@ -160,9 +164,10 @@ class GraphProgramCache { Status & errors) const; mutable std::mutex mutex_; - std::unordered_map> programs_; + std::unordered_map programs_; std::unordered_map validated_matches_; - GraphProgram * last_program_ = nullptr; + std::shared_ptr last_program_; + uint64_t access_tick_ = 0; GraphProgramCacheStats stats_; }; diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index 5ab1fcc193da..8dde8ac42c75 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -126,9 +126,17 @@ std::shared_ptr PreparedCommandProgramCache::g std::lock_guard lock(mutex_); auto found = programs_.find(key); if (found != programs_.end()) { + found->second->last_used = ++access_tick_; return found->second; } + if (programs_.size() >= 32) { + const auto retired = std::min_element( + programs_.begin(), programs_.end(), + [](const auto & a, const auto & b) { return a.second->last_used < b.second->last_used; }); + programs_.erase(retired); + } auto entry = std::make_shared(); + entry->last_used = ++access_tick_; programs_.emplace(key, entry); return entry; } diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index d18ef44375e9..0cb3c1c3488e 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -28,6 +28,7 @@ struct PreparedCommandProgramCacheEntry { PreparedCommandProgram prepared; RecordedCommandGraph recorded; bool ready = false; + uint64_t last_used = 0; }; struct PreparedCommandProgramCacheExecutionResult { @@ -107,6 +108,7 @@ class PreparedCommandProgramCache { mutable std::mutex mutex_; std::unordered_map, KeyHash> programs_; PreparedCommandProgramCacheStats stats_; + uint64_t access_tick_ = 0; }; } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.cpp b/ggml/src/ggml-hrx/runtime/transient-arena.cpp index a12f89e0ab38..a093d58a6b91 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.cpp +++ b/ggml/src/ggml-hrx/runtime/transient-arena.cpp @@ -3,7 +3,9 @@ #include "hrx-interop-utils.h" #include "hrx_runtime.h" +#include #include +#include #include namespace ggml::hrx { @@ -79,15 +81,20 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t status.log("missing HRX stream for transient arena allocation"); return status; } + size_t allocation_size = required_size; + if (!diagnostic_fresh && allocation_capacity_ <= std::numeric_limits::max() / 2) { + // Recorded graphs retain old arenas. Reuse spare capacity across nearby shapes. + allocation_size = std::max(required_size, allocation_capacity_ * 2); + } if (buffer_ != nullptr) { if (diagnostic_fresh) { diagnostic_retired_buffers_.push_back(buffer_); } else { - if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { - status.log("synchronize before growing transient arena: %s", error->c_str()); - return status; - } - hrx_buffer_release(buffer_); + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize before growing transient arena: %s", error->c_str()); + return status; + } + hrx_buffer_release(buffer_); } buffer_ = nullptr; allocation_capacity_ = 0; @@ -101,14 +108,20 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t 0, }; hrx_buffer_t allocation = nullptr; - if (ErrorResult error = take_status( - hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, required_size, &allocation))) { + ErrorResult error = + take_status(hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, allocation_size, &allocation)); + if (error && allocation_size != required_size) { + allocation_size = required_size; + error = take_status( + hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, allocation_size, &allocation)); + } + if (error) { status.log("allocate transient arena: %s", error->c_str()); return status; } buffer_ = allocation; - allocation_capacity_ = required_size; + allocation_capacity_ = allocation_size; allocation_id_ = next_allocation_id(); return status; } diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 292f389b313a..95750faead6c 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -39,6 +39,7 @@ #include #include #include +#include #include #include @@ -2140,8 +2141,8 @@ static void run_unmatched_graph_diagnostics_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 512, 512); - ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, 2); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 7, 512); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 7, 2); ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2); REQUIRE(cache != nullptr); REQUIRE(rows != nullptr); @@ -3282,6 +3283,7 @@ 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 bool q8_vector_kernel = kernel_name == "loom_libs:ggml_mul_mat_vector_q8_q8_1"; const bool aligned_dynamic_kernel = kernel_name == "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"; const bool codebook_kernel = string_contains(kernel_name, "_iq1_") || string_contains(kernel_name, "_iq2_s_") || string_contains(kernel_name, "_iq3_s_") || @@ -3291,9 +3293,9 @@ static void schedule_single_matmul_command(ggml_context * ctx, 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) && q8_capable_kernel; - const size_t command_count = q8_input ? 2 : 1; + const size_t command_count = q8_input || q8_vector_kernel ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); - if (q8_input) { + if (q8_input || q8_vector_kernel) { REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"); } @@ -3311,12 +3313,14 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); if (vector_kernel) { REQUIRE(dispatch.bindings.size() == (codebook_kernel ? 5 : q6_vector_kernel ? 3 : 4)); - if (!q6_vector_kernel) { + if (!q6_vector_kernel && !q8_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 (!q8_vector_kernel) { + require_compile_parameter(dispatch, "ggml.matmul.vector.output_accumulation", "0"); + } if (q6_vector_kernel) { require_compile_parameter(dispatch, "ggml.matmul.vector.q6_storage", "0"); } @@ -5564,6 +5568,8 @@ static void run_qwen_flash_attention_dispatch_checks() { { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, true); + REQUIRE(graph_is_supported(ctx, output)); + output->src[4] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 5); REQUIRE(!graph_is_supported(ctx, output)); } { @@ -5944,6 +5950,9 @@ static void run_qwen_attention_postprocess_dispatch_checks() { REQUIRE(imported.valid()); std::vector covered_nodes(imported.graph.nodes().size(), false); + for (ggml_tensor * input : { tensors.query_raw, tensors.key_raw, tensors.value_raw }) { + covered_nodes[producer_index_for_tensor(imported.graph, input)] = true; + } ggml::hrx::CommandPlan plan; ggml::hrx::DispatchMatch match; REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, @@ -8877,7 +8886,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); + REQUIRE(matmul_graph_is_supported(ctx, output)); } { @@ -8887,7 +8896,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); + REQUIRE(matmul_graph_is_supported(ctx, output)); } ggml_free(ctx); @@ -10394,6 +10403,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + covered_nodes[producer_index_for_tensor(imported.graph, tensors.glu)] = true; ggml::hrx::DispatchMatch down_match = require_common_mul_mat_id_for_graph( imported.graph, plan, covered_nodes, tensors.output, tensors.glu, tensors.down_weight, tensors.route_ids); append_match_to_plan(plan, down_match, covered_nodes, &imported.graph); @@ -10417,6 +10427,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + covered_nodes[producer_index_for_tensor(imported.graph, tensors.glu)] = true; ggml::hrx::DispatchMatch down_match = require_common_mul_mat_id_for_graph( imported.graph, plan, covered_nodes, tensors.output, tensors.glu, tensors.down_weight, tensors.route_ids); append_match_to_plan(plan, down_match, covered_nodes, &imported.graph); @@ -10865,6 +10876,8 @@ static void run_qwen_routed_gate_up_dispatch_checks() { std::vector covered_nodes(imported.graph.nodes().size(), false); const ggml::hrx::CommandPlan empty_plan; ggml::hrx::DispatchMatch gate_up_match; + REQUIRE(!match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, gate_index, gate_up_match)); + covered_nodes[producer_index_for_tensor(imported.graph, tensors.route_ids)] = true; REQUIRE(match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 3); REQUIRE(gate_up_match.transients.size() == 2); @@ -10931,6 +10944,8 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); ggml::hrx::DispatchMatch down_match; + REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + covered_nodes[producer_index_for_tensor(imported.graph, tensors.glu)] = true; 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) == @@ -11096,19 +11111,19 @@ static void run_qwen_router_top8_dispatch_checks() { ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); REQUIRE(logits != nullptr); ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, nullptr, GGML_SORT_ORDER_ASC); - REQUIRE(!graph_is_supported(ctx, output)); + REQUIRE(graph_is_supported(ctx, output)); } { ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, 4); REQUIRE(logits != nullptr); ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits, nullptr, GGML_SORT_ORDER_DESC, 33); - REQUIRE(!graph_is_supported(ctx, output)); + REQUIRE(graph_is_supported(ctx, output)); } { ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 16, 4); REQUIRE(logits != nullptr); ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); - REQUIRE(!graph_is_supported(ctx, output)); + REQUIRE(graph_is_supported(ctx, output)); } { ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kQwenRouterExpertCount, 4); @@ -11667,6 +11682,56 @@ static void run_zero_output_device_support_checks() { ggml_backend_free(backend); } +static void run_generic_hyper_connection_device_support_checks() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + for (int64_t streams : { 1, 4, 65, 128, 129 }) { + auto * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 31, streams, 3); + auto * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, streams, 3); + auto * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 31, 3); + auto * mixing = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, streams, streams, 3); + auto * gate = ggml_dup_tensor(ctx, residual); + REQUIRE(ggml_backend_supports_op(backend, ggml_dsv4_hc_pre(ctx, residual, weights)) == (streams <= 128)); + REQUIRE(ggml_backend_supports_op(backend, ggml_dsv4_hc_pre_gated(ctx, residual, gate, 0.25f)) == + (streams <= 128)); + REQUIRE(ggml_backend_supports_op(backend, ggml_dsv4_hc_post(ctx, input, residual, weights, mixing)) == + (streams <= 128)); + REQUIRE(ggml_backend_supports_op(backend, ggml_dsv4_hc_post(ctx, input, residual, weights, nullptr)) == + (streams <= 128)); + auto * strided = ggml_view_3d(ctx, residual, 15, streams, 3, residual->nb[1], residual->nb[2], 0); + REQUIRE(!ggml_backend_supports_op(backend, ggml_dsv4_hc_pre(ctx, strided, weights))); + } + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_generic_indexer_device_support_checks() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + for (auto [width, heads, supported] : { + std::tuple{ 16, 1, true }, + { 48, 3, true }, + { 128, 63, true }, + { 256, 32, true }, + { 15, 4, false }, + { 128, 65, false }, + { 256, 33, false } + }) { + auto * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, width, heads, 3); + auto * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, width, 1, 17); + auto * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, heads, 3); + auto * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 17, 3); + auto * output = ggml_lightning_indexer(ctx, query, key, weights, mask); + REQUIRE(ggml_backend_supports_op(backend, output) == supported); + } + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_zero_output_unused_weight_staging(bool direct_execution) { EnvironmentVariableGuard replay_mode("GGML_HRX_DISABLE_GRAPH_REPLAY"); if (direct_execution) { @@ -12942,6 +13007,100 @@ static void run_validated_graph_uid_match_checks() { ggml_free(ctx); } +static void run_hyper_connection_pre_cache_checks() { + ggml::hrx::GraphProgramCache cache; + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + auto * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, 1); + auto * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, 1); + auto make_graph = [&](uint64_t uid, bool gated, float scale) { + auto * output = gated ? ggml_dsv4_hc_pre_gated(ctx, x, w, scale) : ggml_dsv4_hc_pre(ctx, x, w); + const auto imported = ggml::hrx::import_op_params(*output); + const auto * params = ggml::hrx::op_params_as(imported); + REQUIRE(params != nullptr && params->gated == gated && params->scale == scale); + auto * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, output); + graph->uid = uid; + return graph; + }; + auto plain = cache.get_or_build(*make_graph(7101, false, 1.0f), corpus, "gfx1151"); + auto gated = cache.get_or_build(*make_graph(7102, true, 1.0f), corpus, "gfx1151"); + auto scaled = cache.get_or_build(*make_graph(7103, true, 0.5f), corpus, "gfx1151"); + auto repeated = cache.get_or_build(*make_graph(7104, true, 1.0f), corpus, "gfx1151"); + REQUIRE(plain.valid() && gated.valid() && scaled.valid() && repeated.valid()); + REQUIRE(plain.program != gated.program); + REQUIRE(gated.program != scaled.program); + REQUIRE(gated.program == repeated.program); + REQUIRE(cache.stats().builds == 3); + ggml_free(ctx); +} + +static void run_swiglu_clamp_cache_checks() { + ggml::hrx::GraphProgramCache cache; + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + auto * gate = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 3); + auto * up = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 3); + auto make_graph = [&](uint64_t uid, float limit) { + auto * output = ggml_swiglu_clamp(ctx, gate, up, limit); + const auto params = ggml::hrx::import_op_params(*output); + REQUIRE(ggml::hrx::op_params_as(params)->limit == limit); + auto * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, output); + graph->uid = uid; + return graph; + }; + auto first = cache.get_or_build(*make_graph(7001, 2.0f), corpus, "gfx1151"); + auto changed = cache.get_or_build(*make_graph(7002, 7.0f), corpus, "gfx1151"); + auto repeated = cache.get_or_build(*make_graph(7003, 2.0f), corpus, "gfx1151"); + REQUIRE(first.valid() && changed.valid() && repeated.valid()); + REQUIRE(first.program != changed.program); + REQUIRE(first.program == repeated.program); + REQUIRE(cache.stats().builds == 2); + ggml_free(ctx); +} + +static void run_graph_program_cache_eviction_checks() { + ggml::hrx::GraphProgramCache cache; + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml_context * ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + auto make_graph = [&](uint64_t uid, int64_t elements) { + auto * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, elements); + auto * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, elements); + auto * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, ggml_add(ctx, a, b)); + graph->uid = uid; + return graph; + }; + auto * first = make_graph(6001, 8); + auto * hot = make_graph(6002, 16); + auto held = cache.get_or_build(*first, corpus, "gfx1151"); + REQUIRE(held.valid()); + std::weak_ptr retired = held.program; + REQUIRE(cache.get_or_build(*hot, corpus, "gfx1151").valid()); + for (int i = 0; i < 64; ++i) { + auto * graph = make_graph(6100 + i, 24 + i * 8); + REQUIRE(cache.get_or_build(*graph, corpus, "gfx1151").valid()); + REQUIRE(cache.get_or_build(*hot, corpus, "gfx1151").valid()); + } + REQUIRE(cache.stats().builds == 66); + REQUIRE(!retired.expired()); + REQUIRE(held.program->match_current_graph(*first).valid()); + held.program.reset(); + REQUIRE(retired.expired()); + auto rebuilt = cache.get_or_build(*first, corpus, "gfx1151"); + REQUIRE(rebuilt.valid()); + REQUIRE(cache.stats().builds == 67); + REQUIRE(cache.get_or_build(*hot, corpus, "gfx1151").valid()); + REQUIRE(cache.stats().builds == 67); + cache.clear(); + REQUIRE(rebuilt.program->match_current_graph(*first).valid()); + ggml_free(ctx); +} + static void run_prepared_graph_instance_reuse() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -12996,7 +13155,7 @@ static void run_prepared_graph_instance_reuse() { auto second = make_graph(4302); auto lookup = state->graph_programs.get_or_build(*first->graph, corpus, state->device->architecture); REQUIRE(lookup.valid()); - auto * program = lookup.program; + auto program = lookup.program; REQUIRE(!program->can_use_prepared_fast_path(*first->graph)); compute(*first, 1); REQUIRE(program->can_use_prepared_fast_path(*first->graph)); @@ -13327,9 +13486,11 @@ static void run_binary_q8_publication_numerics() { legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(test.hidden)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(test.tokens)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", "1"); + legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride0", "1"); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", std::to_string(test.hidden)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", std::to_string(element_count)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", std::to_string(element_count)); + legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride0", "1"); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", std::to_string(test.hidden)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", std::to_string(element_count)); legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", std::to_string(element_count)); @@ -13426,6 +13587,47 @@ static void run_binary_q8_publication_numerics() { ggml_backend_free(backend); } +static void run_transient_arena_growth() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * context = static_cast(backend->context); + ggml::hrx::TransientArena arena; + std::vector retained; + size_t retained_bytes = 0; + uint64_t previous_id = ggml::hrx::kInvalidTransientArenaAllocationId; + constexpr size_t initial = 1024 * 1024; + constexpr size_t maximum = initial + 255 * 4096; + REQUIRE(arena.ensure_capacity(context->device->device, context->stream, 0).success()); + REQUIRE(arena.capacity() == 0); + for (size_t step = 0; step < 256; ++step) { + const size_t required = initial + step * 4096; + REQUIRE(arena.ensure_capacity(context->device->device, context->stream, required).success()); + const auto allocation = arena.current_allocation(); + REQUIRE(allocation.buffer != nullptr); + REQUIRE(allocation.capacity >= required); + REQUIRE(allocation.allocation_id != ggml::hrx::kInvalidTransientArenaAllocationId); + if (allocation.allocation_id != previous_id) { + hrx_buffer_retain(allocation.buffer); + retained.push_back(allocation); + retained_bytes += allocation.capacity; + previous_id = allocation.allocation_id; + } + } + REQUIRE(arena.ensure_capacity(context->device->device, context->stream, initial).success()); + REQUIRE(arena.allocation_id() == previous_id); + arena.clear(); + REQUIRE(arena.capacity() == 0); + REQUIRE(arena.allocation_id() == ggml::hrx::kInvalidTransientArenaAllocationId); + for (const auto & allocation : retained) { + hrx_buffer_release(allocation.buffer); + } + std::fprintf(stderr, "transient growth: %zu allocations, %zu retained bytes, %zu required bytes\n", retained.size(), + retained_bytes, maximum); + ggml_backend_free(backend); + REQUIRE(retained.size() <= 8); + REQUIRE(retained_bytes <= 4 * maximum); +} + static void run_qwen_expert_table_partition_prefill_512_execution() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -13956,6 +14158,66 @@ static void run_two_independent_add_f32() { ggml_backend_free(backend); } +static void run_prepared_program_cache_eviction() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_context * ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + + struct Case { + ggml_tensor * a; + ggml_tensor * b; + ggml_tensor * out; + ggml_cgraph * graph; + }; + + std::vector cases; + for (int i = 0; i < 65; ++i) { + auto * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + auto * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + auto * out = ggml_add(ctx, a, b); + auto * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, out); + graph->uid = 7001; + cases.push_back({ a, b, out, graph }); + } + auto buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + auto run = [&](size_t index, float base) { + auto & test = cases[index]; + std::array a, b, actual; + for (size_t i = 0; i < a.size(); ++i) { + a[i] = base + float(i) * .25f; + b[i] = -base * .5f + float(i) * .125f; + } + ggml_backend_tensor_set(test.a, a.data(), 0, sizeof(a)); + ggml_backend_tensor_set(test.b, b.data(), 0, sizeof(b)); + REQUIRE(ggml_backend_graph_compute(backend, test.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(test.out, actual.data(), 0, sizeof(actual)); + for (size_t i = 0; i < actual.size(); ++i) { + REQUIRE(actual[i] == a[i] + b[i]); + } + }; + for (size_t i = 0; i < cases.size(); ++i) { + run(i, float(i) - 17.0f); + } + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.graph_program_builds == 1); + REQUIRE(stats.prepared_program_builds == 65); + run(1, 91.0f); + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.prepared_program_builds == 66); + run(64, -63.0f); + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.prepared_program_builds == 66); + REQUIRE(stats.prepared_program_hits >= 1); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_chained_add_f32() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -14661,6 +14923,9 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); suite.host_case("graph_compute_mask", [] { run_graph_compute_mask_checks(); }); suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); + suite.host_case("graph_program_cache_eviction", [] { run_graph_program_cache_eviction_checks(); }); + suite.host_case("swiglu_clamp_cache", [] { run_swiglu_clamp_cache_checks(); }); + suite.host_case("hyper_connection_pre_cache", [] { run_hyper_connection_pre_cache_checks(); }); suite.host_case("graph_match_empty_expert_views", [] { run_graph_match_empty_expert_views_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(); }); @@ -14800,12 +15065,15 @@ static void register_hrx_backend_device_cases(test_runner::Suite & suite) { run_registered_host_input_reuse_waits_before_mutation(); }); suite.device_case("initializer_host_input_ownership", [] { run_initializer_host_input_ownership(); }); - + suite.device_case("transient_arena_growth", [] { run_transient_arena_growth(); }); suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); suite.device_case("zero_output_unused_weight_staging.recorded", [] { run_zero_output_unused_weight_staging(false); }); suite.device_case("zero_output_unused_weight_staging.direct", [] { run_zero_output_unused_weight_staging(true); }); suite.device_case("prepared_graph_instance_reuse", [] { run_prepared_graph_instance_reuse(); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); + suite.device_case("generic_hyper_connection_device_support", + [] { run_generic_hyper_connection_device_support_checks(); }); + suite.device_case("generic_indexer_device_support", [] { run_generic_indexer_device_support_checks(); }); suite.device_case("binary_q8_publication_numerics", [] { run_binary_q8_publication_numerics(); }); suite.device_case("qwen_expert_table_partition_prefill_512_execution", [] { run_qwen_expert_table_partition_prefill_512_execution(); }); @@ -14821,6 +15089,7 @@ static void register_hrx_backend_device_cases(test_runner::Suite & suite) { [] { 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("prepared_program_cache_eviction", [] { run_prepared_program_cache_eviction(); }); suite.device_case("unsupported_op_fails", [] { run_unsupported_op_fails(); }); } diff --git a/tests/hrx-loom-jit-test.cpp b/tests/hrx-loom-jit-test.cpp index c335d1a08c86..3a9609f0b2f9 100644 --- a/tests/hrx-loom-jit-test.cpp +++ b/tests/hrx-loom-jit-test.cpp @@ -134,18 +134,20 @@ static const ggml::hrx::KernelDefinition * find_targeted_kernel(const char * nam static std::map binary_f32_exact_config(const char * op, int64_t element_count) { const std::string count = std::to_string(element_count); return { - { "ggml.binary_f32.op", op }, - { "ggml.binary_f32.ne0", count }, - { "ggml.binary_f32.ne1", "1" }, - { "ggml.binary_f32.ne2", "1" }, + { "ggml.binary_f32.op", op }, + { "ggml.binary_f32.ne0", count }, + { "ggml.binary_f32.ne1", "1" }, + { "ggml.binary_f32.ne2", "1" }, + { "ggml.binary_f32.src0_stride0", "1" }, { "ggml.binary_f32.src0_stride1", count }, { "ggml.binary_f32.src0_stride2", count }, { "ggml.binary_f32.src0_stride3", count }, + { "ggml.binary_f32.src1_stride0", "1" }, { "ggml.binary_f32.src1_stride1", count }, { "ggml.binary_f32.src1_stride2", count }, { "ggml.binary_f32.src1_stride3", count }, - { "ggml.binary_f32.src0_span", count }, - { "ggml.binary_f32.src1_span", count }, + { "ggml.binary_f32.src0_span", count }, + { "ggml.binary_f32.src1_span", count }, }; } diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index feed4d3ed733..b0019f753206 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -1906,7 +1906,9 @@ static void run_gemma_post_proj_matcher_negatives() { imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, plan, next_plan_value(imported.graph, plan), }; - REQUIRE(registry->match(context, match)); + if (!registry->match(context, match)) { + return std::string(); + } REQUIRE(match.dispatches.size() == 1); const auto & specialization = match.dispatches[0].kernel; const std::string name = kernel_name_for_id(specialization.kernel_id); @@ -1925,7 +1927,7 @@ static void run_gemma_post_proj_matcher_negatives() { } return name; }; - REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); + REQUIRE(kernel().empty()); covered[raw_index] = true; if (variant == 3) { REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); @@ -4577,12 +4579,14 @@ static void run_qwen_flash_attention_case() { ggml_backend_free(backend); } -static void run_common_flash_attention_cpu_reference_case(int64_t head_size, - int64_t query_token_count, - int64_t key_value_token_count, - float scale, - const char * expected_kernel, - int64_t value_head_size = -1) { +static void run_common_flash_attention_cpu_reference_case( + int64_t head_size, + int64_t query_token_count, + int64_t key_value_token_count, + float scale, + const char * expected_kernel, + int64_t value_head_size = -1, + ggml::hrx::WorkloadSpecialization expected_specialization = ggml::hrx::WorkloadSpecialization::Dynamic) { if (value_head_size < 0) { value_head_size = head_size; } @@ -4616,9 +4620,6 @@ static void run_common_flash_attention_cpu_reference_case(int64_t head_size ggml_build_forward_expand(graph, output); require_kernel_subsequence(scheduled_kernel_sequence(graph), { expected_kernel }); - const ggml::hrx::WorkloadSpecialization expected_specialization = - query_token_count == 256 && key_value_token_count == 256 ? ggml::hrx::WorkloadSpecialization::Exact : - ggml::hrx::WorkloadSpecialization::Dynamic; REQUIRE(scheduled_kernel_specialization(graph, expected_kernel).workload_specialization == expected_specialization); @@ -4654,6 +4655,105 @@ static void run_common_flash_attention_cpu_reference_case(int64_t head_size restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); } +static void run_flash_attention_masked_storage_case(int64_t query_count, int64_t key_count, int64_t visible_count) { + constexpr int64_t head_size = 512; + const int64_t head_count = query_count == 1 ? 64 : 4; + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_context * ctx = ggml_init({ 4 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, query_count, head_count); + query->nb[1] = head_count * head_size * sizeof(float); + query->nb[2] = head_size * sizeof(float); + ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_size, key_count, 1); + ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_size, key_count, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_count, query_count); + ggml_tensor * sink = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_count); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask, head_size); + ggml_flash_attn_ext_add_sinks(output, sink); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { "loom_libs:ggml_flash_attention_f32_f16_wmma_sink" }); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + uint32_t random_state = 20261028; + auto random_value = [&]() { + random_state ^= random_state << 13; + random_state ^= random_state >> 17; + random_state ^= random_state << 5; + return (static_cast(random_state & 65535) / 65536.0f - 0.5f) * 8.0f; + }; + std::vector query_data(query_count * head_count * head_size, 0.0f); + std::vector sink_data(head_count, 0.0f); + std::vector key_data(key_count * head_size), value_data(key_count * head_size); + for (size_t i = 0; i < key_data.size(); ++i) { + key_data[i] = ggml_fp32_to_fp16(random_value()); + value_data[i] = ggml_fp32_to_fp16(random_value()); + } + std::vector mask_data(query_count * key_count, + ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + std::vector expected(query_data.size(), 0.0f); + for (int64_t row = 0; row < query_count; ++row) { + const int64_t visible = std::max(0, visible_count - query_count + 1 + row); + for (int64_t token = 0; token < visible; ++token) { + mask_data[row * key_count + token] = ggml_fp32_to_fp16(0.0f); + } + // Zero queries and sinks give uniform probabilities, including one sink entry. + for (int64_t channel = 0; channel < head_size; ++channel) { + double sum = 0.0; + for (int64_t token = 0; token < visible; ++token) { + sum += ggml_fp16_to_fp32(value_data[token * head_size + channel]); + } + for (int64_t head = 0; head < head_count; ++head) { + expected[(row * head_count + head) * head_size + channel] = sum / (visible + 1); + } + } + } + ggml_backend_tensor_set(query, query_data.data(), 0, ggml_nbytes(query)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, ggml_nbytes(mask)); + ggml_backend_tensor_set(sink, sink_data.data(), 0, ggml_nbytes(sink)); + std::vector baseline(expected.size()), actual(expected.size()); + for (int variant = 0; variant < 5; ++variant) { + auto changed_key = key_data; + auto changed_value = value_data; + for (int64_t i = visible_count * head_size; i < key_count * head_size; ++i) { + for (auto * data : { &changed_key, &changed_value }) { + const float magnitude = std::fabs(ggml_fp16_to_fp32((*data)[i])); + if (variant > 0) { + const float replacement = variant == 1 ? magnitude : + variant == 2 ? -magnitude : + variant == 3 ? 0.0f : + -0.0f; + (*data)[i] = ggml_fp32_to_fp16(replacement); + } + } + } + ggml_backend_tensor_set(key, changed_key.data(), 0, ggml_nbytes(key)); + ggml_backend_tensor_set(value, changed_value.data(), 0, ggml_nbytes(value)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(output, actual.data(), 0, ggml_nbytes(output)); + double squared_error = 0.0, squared_reference = 0.0; + for (size_t i = 0; i < actual.size(); ++i) { + REQUIRE(std::isfinite(actual[i])); + const double error = static_cast(actual[i]) - expected[i]; + squared_error += error * error; + squared_reference += static_cast(expected[i]) * expected[i]; + } + REQUIRE(squared_error <= 1e-5 * squared_reference); + if (variant == 0) { + baseline = actual; + } else { + REQUIRE(std::memcmp(actual.data(), baseline.data(), actual.size() * sizeof(float)) == 0); + } + } + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_asymmetric_flash_attention_cpu_reference_case() { static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); @@ -4691,8 +4791,9 @@ static void run_asymmetric_flash_attention_cpu_reference_case() { ggml_build_forward_expand(graph, output); require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_flash_attention_f32_f16_wmma" }); - REQUIRE(scheduled_kernel_specialization(graph, "loom_libs:ggml_flash_attention_f32_f16_wmma") - .workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE( + scheduled_kernel_specialization(graph, "loom_libs:ggml_flash_attention_f32_f16_wmma").workload_specialization == + ggml::hrx::WorkloadSpecialization::Exact); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -4758,6 +4859,12 @@ static void run_qwen_full_cache_prefill_flash_attention_scheduling_case(int64_t std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan; ggml::hrx::DispatchMatch postprocess_match; + // This fixture starts after the three input projections have completed. + for (size_t i = 0; i < imported.graph.nodes().size(); ++i) { + if (imported.graph.nodes()[i].op == GGML_OP_MUL_MAT) { + covered_nodes[i] = true; + } + } match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, producer_index_for_tensor(imported.graph, attention.query_reshape), postprocess_match); ggml::hrx::DispatchMatch flash_match; @@ -7939,6 +8046,8 @@ static const char * glu_op_name(ggml_glu_op op) { return "swiglu"; case GGML_GLU_OP_SWIGLU_OAI: return "swiglu_oai"; + case GGML_GLU_OP_SWIGLU_CLAMP: + return "swiglu_clamp"; case GGML_GLU_OP_GEGLU_ERF: return "geglu_erf"; case GGML_GLU_OP_GEGLU_QUICK: @@ -8499,6 +8608,7 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { 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(); }); + for (const int64_t head_count : { 8, 32 }) { for (const int64_t token_count : { 1, 64 }) { suite.device_case( @@ -8569,9 +8679,23 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { } static void register_flash_attention_cases(Suite & suite) { + for (const auto & shape : std::vector>{ + { 1, 31, 23 }, + { 1, 64, 7 }, + { 1, 65, 63 }, + { 1, 256, 7 }, + { 17, 200, 197 }, + { 4, 257, 253 }, + { 1, 31, 0 } + }) { + const std::string name = "common_flash_attention.masked_storage.query" + std::to_string(shape[0]) + ".kv" + + std::to_string(shape[1]) + ".visible" + std::to_string(shape[2]); + suite.device_case(name, [shape] { run_flash_attention_masked_storage_case(shape[0], shape[1], shape[2]); }); + } suite.device_case("common_flash_attention.prefill.head128.query256.kv256.exact", [] { run_common_flash_attention_cpu_reference_case(128, 256, 256, 1.0f / std::sqrt(128.0f), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); + "loom_libs:ggml_flash_attention_f32_f16_wmma", -1, + ggml::hrx::WorkloadSpecialization::Exact); }); for (int64_t head_size = 64; head_size <= 512; head_size += 64) { suite.device_case("common_flash_attention.decode.head" + std::to_string(head_size) + ".query2.kv65", @@ -8582,9 +8706,9 @@ static void register_flash_attention_cases(Suite & suite) { }); 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"); + 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", -1, ggml::hrx::WorkloadSpecialization::Exact); }); suite.device_case( "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query17.kv65", [head_size] { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 302851fdeb29..c0c7c00c2196 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2611,6 +2611,30 @@ struct test_mul_mat_reorder_then_big_batch : public test_case { } }; +struct test_get_rows_router : public test_get_rows { + using test_get_rows::test_get_rows; + + std::string vars() override { return test_get_rows::vars() + ",router_gather=1"; } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->op != GGML_OP_NONE) { + continue; + } + if (t->type == GGML_TYPE_I32) { + std::vector data(ggml_nelements(t)); + const bool ids = std::string(t->name) == "rows"; + for (size_t i = 0; i < data.size(); ++i) { + data[i] = ids ? (i * 31 + 7) % m : i * 37 + 19; + } + ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); + } else { + init_tensor_uniform(t); + } + } + } +}; + // GGML_OP_GET_ROWS_BACK struct test_get_rows_back : public test_case { const ggml_type type; @@ -3605,6 +3629,27 @@ static void mask_skipped_ids(std::vector & ids, int64_t row) { } } +struct test_bin_strided_bcast : public test_case { + test_bin_bcast::op_t op; + int64_t width, tokens; + bool scalar; + + test_bin_strided_bcast(test_bin_bcast::op_t op, int64_t width, int64_t tokens, bool scalar) : + op(op), + width(width), + tokens(tokens), + scalar(scalar) {} + + std::string vars() override { return VARS_TO_STR3(width, tokens, scalar) + ",row_stride=24"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 24, tokens); + ggml_tensor * lhs = ggml_view_2d(ctx, storage, width, tokens, storage->nb[1], 4 * sizeof(float)); + ggml_tensor * rhs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scalar ? 1 : width); + return op(ctx, lhs, rhs); + } +}; + // GGML_OP_ADD_ID struct test_add_id : public test_case { const ggml_type type_a; @@ -4604,24 +4649,26 @@ struct test_dsv4_hc_comb : public test_dsv4_hc { const int64_t n_tokens; const int32_t n_iter; const float eps; + const int64_t n_scale; std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); return "DSV4_HC_COMB"; } - std::string vars() override { - return VARS_TO_STR3(n_tokens, n_iter, eps); - } + std::string vars() override { return VARS_TO_STR4(n_tokens, n_iter, eps, n_scale); } - test_dsv4_hc_comb(int64_t n_tokens = 17, int32_t n_iter = 4, float eps = 1e-6f) - : n_tokens(n_tokens), n_iter(n_iter), eps(eps) {} + test_dsv4_hc_comb(int64_t n_tokens = 17, int32_t n_iter = 4, float eps = 1e-6f, int64_t n_scale = 3) : + n_tokens(n_tokens), + n_iter(n_iter), + eps(eps), + n_scale(n_scale) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * mixes = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (2 + hc)*hc, n_tokens); ggml_set_name(mixes, "mixes"); - ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 3); + ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_scale); ggml_set_name(scale, "scale"); ggml_tensor * base = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (2 + hc)*hc); @@ -4676,39 +4723,46 @@ struct test_dsv4_hc_post : public test_dsv4_hc { const int64_t n_tokens; const bool identity; const bool gated; + const int64_t n_hc; std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); return "DSV4_HC_POST"; } - std::string vars() override { - return VARS_TO_STR4(n_embd, n_tokens, identity, gated); - } + std::string vars() override { return VARS_TO_STR5(n_embd, n_tokens, identity, gated, n_hc); } // gated: post = 2*sigmoid(post/hc), as qwen4exp builds it, so backends can fuse the chain bool run_whole_graph() override { return gated; } - test_dsv4_hc_post(int64_t n_embd = 31, int64_t n_tokens = 17, bool identity = false, bool gated = false) - : n_embd(n_embd), n_tokens(n_tokens), identity(identity), gated(gated) {} + test_dsv4_hc_post(int64_t n_embd = 31, + int64_t n_tokens = 17, + bool identity = false, + bool gated = false, + int64_t n_hc = 4) : + n_embd(n_embd), + n_tokens(n_tokens), + identity(identity), + gated(gated), + n_hc(n_hc) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); ggml_set_name(x, "x"); - ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, hc, n_tokens); + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_hc, n_tokens); ggml_set_name(residual, "residual"); - ggml_tensor * post = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc, n_tokens); + ggml_tensor * post = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_hc, n_tokens); ggml_set_name(post, "post"); if (gated) { - post = ggml_scale(ctx, ggml_sigmoid(ctx, ggml_scale(ctx, post, 1.0f / (float) hc)), 2.0f); + post = ggml_scale(ctx, ggml_sigmoid(ctx, ggml_scale(ctx, post, 1.0f / (float) n_hc)), 2.0f); } ggml_tensor * comb = nullptr; if (!identity) { - comb = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hc, hc, n_tokens); + comb = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_hc, n_hc, n_tokens); ggml_set_name(comb, "comb"); } @@ -5540,6 +5594,68 @@ struct test_mul_mat : public test_case { } }; +struct test_mul_mat_late_bias : public test_mul_mat { + const bool bias_first; + + test_mul_mat_late_bias(bool bias_first) : + test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 256, 17, 256, { 1, 1 }, { 1, 1 }), bias_first(bias_first) {} + + bool run_whole_graph() override { return true; } + + std::string vars() override { return test_mul_mat::vars() + ",late_bias=1," + VAR_TO_STR(bias_first); } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * projection = test_mul_mat::build_graph(ctx); + ggml_tensor * bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, m, n); + bias = ggml_scale(ctx, bias, 0.25f); + return bias_first ? ggml_add(ctx, bias, projection) : ggml_add(ctx, projection, bias); + } +}; + +struct test_mul_mat_small_k : public test_mul_mat { + using test_mul_mat::test_mul_mat; + + std::string vars() override { return test_mul_mat::vars() + ",small_k=1"; } +}; + +struct test_mul_mat_long_positive : public test_mul_mat { + test_mul_mat_long_positive(int64_t tokens) : + test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 24, tokens, 16384, { 1, 1 }, { 1, 1 }) {} + + std::string vars() override { return test_mul_mat::vars() + ",long_positive=1"; } + + double max_nmse_err() override { return 1e-7; } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (ggml_is_view_op(t->op)) { + continue; + } + init_tensor_uniform(t, 0.05f, 0.1f); + } + } +}; + +struct test_mul_mat_grouped : public test_mul_mat { + const bool strided_input; + + test_mul_mat_grouped(ggml_type type, int64_t tokens, bool strided_input) : + test_mul_mat(type, GGML_TYPE_F32, 1024, tokens, 4096, { 8, 1 }, { 1, 1 }), + strided_input(strided_input) {} + + std::string vars() override { return test_mul_mat::vars() + ",grouped_projection=1," + VAR_TO_STR(strided_input); } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * weight = ggml_new_tensor_3d(ctx, type_a, k, m, bs[0]); + ggml_tensor * input = + strided_input ? ggml_new_tensor_3d(ctx, type_b, k, bs[0], n) : ggml_new_tensor_3d(ctx, type_b, k, n, bs[0]); + if (strided_input) { + input = ggml_permute(ctx, input, 0, 2, 1, 3); + } + return ggml_mul_mat(ctx, weight, input); + } +}; + // GGML_HINT_SRC0_IS_HADAMARD struct test_mul_mat_hadamard : public test_mul_mat { test_mul_mat_hadamard(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, @@ -8782,6 +8898,28 @@ struct test_flash_attn_ext_large_logits : public test_flash_attn_ext { } }; +struct test_flash_attn_sink : public test_flash_attn_ext { + explicit test_flash_attn_sink(int64_t tokens, int64_t head_size = 512, int64_t heads = 64, int64_t keys = 256) : + test_flash_attn_ext(head_size, head_size, 1, { heads, 1 }, keys, tokens, true, true) {} + + std::string vars() override { return test_flash_attn_ext::vars() + ",model_layout=1"; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + auto * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hsk, nh * nr23[0], nb); + q = ggml_permute(ctx, q, 0, 2, 1, 3); + auto * k = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hsk, kv); + auto * v = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hsv, kv); + auto * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); + auto * s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nh * nr23[0]); + ggml_set_name(m, "m"); + ggml_set_name(s, "s"); + auto * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f / sqrtf(hsk), 0.0f, 0.0f); + ggml_flash_attn_ext_add_sinks(out, s); + ggml_prec_set_acc(out, GGML_PREC_F32); + return out; + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -10107,10 +10245,20 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_snake_fuse(type, { 64, 32, 2, 3})); // ne[2] > 1 and ne[3] > 1 } + for (auto op : { ggml_add, ggml_sub, ggml_mul, ggml_div }) { + for (int64_t tokens : { 1, 17, 65 }) { + for (bool scalar : { false, true }) { + test_cases.emplace_back(new test_bin_strided_bcast(op, 4, tokens, scalar)); + test_cases.emplace_back(new test_bin_strided_bcast(op, 7, tokens, scalar)); + } + } + } + test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 20)); + test_cases.emplace_back(new test_dsv4_hc_comb(17, 4, 1e-6f, 7)); // production n_iter (DeepSeek-V4 uses 20) across batch sizes that cross // subgroup and workgroup boundaries; 1 = single-token decode for (int64_t n_tokens : {1, 256, 336, 512, 513, 1024, 2048}) { @@ -10123,7 +10271,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, 21)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 4, 17, true)); test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, 21, true)); - for (int64_t n_hc : {1, 2, 3, 5, 8, 65}) { + for (int64_t n_hc : { 1, 2, 3, 5, 8, 65, 128 }) { test_cases.emplace_back(new test_dsv4_hc_pre(128, n_hc, 17)); test_cases.emplace_back(new test_dsv4_hc_pre(128, n_hc, 17, true)); } @@ -10137,6 +10285,18 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_post(31, 17, true, true)); test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true)); test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true)); + for (int64_t n_hc : { 1, 3, 8, 65, 128 }) { + test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, false, n_hc)); + test_cases.emplace_back(new test_dsv4_hc_post(31, 17, true, false, n_hc)); + } + for (int64_t tokens : { 1, 2, 3, 4, 5, 17, 512 }) { + test_cases.emplace_back(new test_flash_attn_sink(tokens)); + } + for (int64_t tokens : { 1, 2, 3, 4 }) { + test_cases.emplace_back(new test_flash_attn_sink(tokens, 64, 4, 63)); + test_cases.emplace_back(new test_flash_attn_sink(tokens, 128, 12, 97)); + test_cases.emplace_back(new test_flash_attn_sink(tokens, 256, 32, 129)); + } for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { // blocks with extra cells and scores (k-pool) @@ -10197,6 +10357,12 @@ static std::vector> make_test_cases_eval() { } } + for (int tokens : { 1, 4, 512 }) { + for (float limit : { 0.0f, 7.0f, 10.0f }) { + test_cases.emplace_back(new test_swiglu_clamp(GGML_TYPE_F32, { 2048, 6, tokens, 1 }, 0, limit)); + } + } + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_0}) { test_cases.emplace_back(new test_get_rows(type, 300*256, 5, 4, 1, 2, false)); test_cases.emplace_back(new test_get_rows(type, 256, 80000, 70000, 2, 1, false)); @@ -10211,6 +10377,11 @@ static std::vector> make_test_cases_eval() { for (ggml_type type : {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S}) { test_cases.emplace_back(new test_mul_mat_reorder_then_big_batch(type)); } + for (int tokens : { 1, 17, 512 }) { + test_cases.emplace_back(new test_get_rows_router(GGML_TYPE_I32, 6, 129280, tokens, 1, 1, false)); + test_cases.emplace_back(new test_get_rows_router(GGML_TYPE_F32, 1, 256, 6, tokens, 1, false)); + test_cases.emplace_back(new test_get_rows_router(GGML_TYPE_F32, 1, 256, 12, tokens, 2, true)); + } for (ggml_type type : all_types) { for (int b : {1, 7}) { for (bool v : {false, true}) { @@ -10242,6 +10413,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + for (ggml_type cache : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + for (int tokens : { 1, 17, 512 }) { + test_cases.emplace_back( + new test_set_rows(GGML_TYPE_F32, cache, GGML_TYPE_I32, { 1024, 600, 1, 1 }, { 1, 1 }, tokens)); + } + } test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false)); for (ggml_type src_type : {GGML_TYPE_F16, GGML_TYPE_F32}) { for (ggml_type type : all_types) { @@ -10802,6 +10979,8 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_cont(GGML_TYPE_F32, { 512, 8, 3, 1 })); + test_cases.emplace_back(new test_cont(GGML_TYPE_F32, { 512, 128, 4, 1 })); auto add_test_bin_bcast = [&](ggml_type type, std::array ne, std::array nr, bool perm1 = false, bool src_overlap = false) { for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) { @@ -10887,6 +11066,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 1}, 16)); test_cases.emplace_back(new test_scale()); + test_cases.emplace_back(new test_scale(GGML_TYPE_F32, { 1024, 1, 1, 1 }, 0.0f, -INFINITY)); + test_cases.emplace_back(new test_scale(GGML_TYPE_F32, { 257, 1, 1, 1 }, 1.0f, INFINITY)); test_cases.emplace_back(new test_scale(GGML_TYPE_F32, {10, 10, 10, 10}, 2.0f, 1.0f)); test_cases.emplace_back(new test_scale(GGML_TYPE_F32, {10, 10, 10, 10}, 2.0f, 1.0f, true)); // inplace test test_cases.emplace_back(new test_scale(GGML_TYPE_F32, {100, 10, 10, 10}, 2.0f, 1.0f)); @@ -11480,6 +11661,46 @@ static std::vector> make_test_cases_eval() { // this case is verified (pass) in Intel(R) Data Center GPU Max 1100 (sycl backend) and NV A30 (cuda backend) // test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 512, 262144, 9216, {1, 1}, {1, 1})); + for (ggml_type type : { GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_F32 }) { + for (int tokens : { 1, 17, 512 }) { + for (bool strided : { false, true }) { + test_cases.emplace_back(new test_mul_mat_grouped(type, tokens, strided)); + } + } + } + + for (bool bias_first : { false, true }) { + test_cases.emplace_back(new test_mul_mat_late_bias(bias_first)); + } + for (int tokens : { 1, 17, 512 }) { + test_cases.emplace_back(new test_rope(GGML_TYPE_F32, { 64, 64, tokens, 1 }, 64, GGML_ROPE_TYPE_NORMAL, 1024, + 0.0625f, 1.0f, 1.0f, false, 1, false)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, { 256, tokens, 1, 1 }, 256, true)); + } + for (int tokens : { 1, 8, 512 }) { + test_cases.emplace_back(new test_mul_mat_long_positive(tokens)); + test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F16, { 256, tokens, 1, 1 }, { 1, 1, 1, 1 })); + } + for (ggml_type type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { + for (ggml_type ids : { GGML_TYPE_I32, GGML_TYPE_I64 }) { + for (int tokens : { 1, 17, 512 }) { + test_cases.emplace_back(new test_set_rows(type, type, ids, { 1, 256, tokens, 1 }, { 1, 1 }, 64)); + } + } + } + for (ggml_type type : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + for (int tokens : { 1, 17, 512 }) { + for (bool strided : { false, true }) { + const std::array per = + strided ? std::array{ 0, 2, 1, 3 } : std::array{ 0, 1, 2, 3 }; + test_cases.emplace_back( + new test_mul_mat_small_k(type, GGML_TYPE_F32, 128, tokens, 128, { 1, 1 }, { 1, 1 }, per)); + test_cases.emplace_back( + new test_mul_mat_small_k(type, GGML_TYPE_F32, 256, tokens, 128, { 1, 1 }, { 64, 1 }, per)); + } + } + } + // test large experts*tokens for (bool b : {false, true}) { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, b, 32, 1024, 16)); @@ -11512,6 +11733,17 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_S, GGML_TYPE_F32, 1024, 10, false, 128, n, 512)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1024, 10, false, 256, n, 128)); } + for (bool broadcast : { false, true }) { + for (int k : { 512, 2048, 4096 }) { + for (int n : { 1, 2, 3, 4, 5 }) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 6, broadcast, 65, n, k)); + } + } + for (int n : { 1, 2, 3, 4, 5 }) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 6, broadcast, 1025, n, 2048, 0, n == 3, n == 5)); + } + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 6, broadcast, 2048, 4, 4096)); + } test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880)); // multiple blocks per row: exercises the block-stride loop and the @@ -11733,6 +11965,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {16, 2, 32, 1}, mask, sinks, GGML_TYPE_F16, {1, 1}, 0.1f, 0.0f, true)); } } + test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, { 8, 512, 3, 1 })); + test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, { 128, 512, 4, 1 })); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {16, 2, 32, 1}, true, true, GGML_TYPE_F32, {1, 1}, 0.1f, 0.0f)); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {16, 2, 32, 1}, true, false, GGML_TYPE_F16, {1, 1}, 0.1f, 0.0f)); test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {16, 2, 32, 1}, false, true, GGML_TYPE_F32, {1, 1}, 0.1f, 0.0f)); @@ -12008,6 +12242,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_sum(GGML_TYPE_F32, {11, 5, 6, 3}, {0, 0, 0, 0}, true)); // sliced view (gaps) test_cases.emplace_back(new test_sum(GGML_TYPE_F32, {11, 5, 6, 3}, {1, 0, 2, 3}, true)); // sliced + non-contiguous rows test_cases.emplace_back(new test_sum_rows()); + for (int tokens : { 1, 17, 512 }) { + test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 6, tokens, 1, 1 })); + } test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, false)); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, false, true)); test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, true)); @@ -12087,6 +12324,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_tri(GGML_TRI_TYPE_UPPER_DIAG)); test_cases.emplace_back(new test_fill(0.0f)); + for (ggml_type type : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + test_cases.emplace_back(new test_fill(-INFINITY, type, { 256, 512, 1, 1 })); + test_cases.emplace_back(new test_fill(0.0f, type, { 1, 256, 512, 1 })); + } test_cases.emplace_back(new test_fill(2.0f, GGML_TYPE_F32, { 303, 207, 11, 3 })); test_cases.emplace_back(new test_fill(-152.0f, GGML_TYPE_F32, { 800, 600, 4, 4 })); test_cases.emplace_back(new test_fill(3.5f, GGML_TYPE_F32, { 2048, 512, 2, 2 })); @@ -12566,6 +12807,20 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_falcon(2)); #endif + // Cover head tiles, width tiles, and partial key tiles independently. + for (auto [head_size, heads] : { + std::pair{ 16, 1 }, + { 48, 3 }, + { 64, 15 }, + { 64, 16 }, + { 128, 17 }, + { 256, 32 }, + { 128, 48 }, + { 128, 63 } + }) { + test_cases.emplace_back(new test_lightning_indexer(head_size, heads, 17, 3)); + } + // lightning_indexer for (int kv : { 256 }) { for (int bs : { 1, 512 }) {