diff --git a/ggml/src/ggml-hrx/CMakeLists.txt b/ggml/src/ggml-hrx/CMakeLists.txt index 43e9e7cf67e9..b3eb70120f57 100644 --- a/ggml/src/ggml-hrx/CMakeLists.txt +++ b/ggml/src/ggml-hrx/CMakeLists.txt @@ -199,6 +199,8 @@ ggml_add_backend_library(ggml-hrx dispatch/dispatch.h dispatch/transient-allocator.cpp dispatch/transient-allocator.h + dispatch_registration/common/dispatch-activation-publication.cpp + dispatch_registration/common/dispatch-activation-publication.h dispatch_registration/common/dispatch-binary.cpp dispatch_registration/common/dispatch-binary.h dispatch_registration/common/dispatch-binary-common.h @@ -234,8 +236,12 @@ ggml_add_backend_library(ggml-hrx dispatch_registration/common/dispatch-scale.h dispatch_registration/common/dispatch-unary.cpp dispatch_registration/common/dispatch-unary.h + dispatch_registration/common/iq-codebook-resource.cpp + dispatch_registration/common/iq-codebook-resource.h dispatch_registration/dispatch-registry.cpp dispatch_registration/dispatch-registry.h + dispatch_registration/immutable-program-resource.cpp + dispatch_registration/immutable-program-resource.h dispatch_registration/llm/dispatch-attention-qkv.cpp dispatch_registration/llm/dispatch-attention-qkv.h dispatch_registration/llm/dispatch-gated-delta-net.cpp diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp index 89d9572d66f8..96e3faaa202f 100644 --- a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp @@ -37,6 +37,7 @@ static bool moe_routing_bundle_matches(const CommandPlanMoeRoutingBundle & lhs, void CommandPlanMetadata::clear() { generated_resources_.clear(); alternate_values_.clear(); + activation_publication_diagnostics_.clear(); moe_routing_bundles_.clear(); } @@ -51,6 +52,9 @@ bool CommandPlanMetadata::append(CommandPlanMetadata && other, Status & status) return false; } } + for (CommandPlanActivationPublicationDiagnostic & diagnostic : other.activation_publication_diagnostics_) { + append_activation_publication_diagnostic(std::move(diagnostic)); + } for (CommandPlanMoeRoutingBundle & bundle : other.moe_routing_bundles_) { if (!append_moe_routing_bundle(std::move(bundle), status)) { return false; @@ -59,6 +63,11 @@ bool CommandPlanMetadata::append(CommandPlanMetadata && other, Status & status) return true; } +void CommandPlanMetadata::append_activation_publication_diagnostic( + CommandPlanActivationPublicationDiagnostic diagnostic) { + activation_publication_diagnostics_.push_back(std::move(diagnostic)); +} + bool CommandPlanMetadata::append_generated_resource(CommandPlanGeneratedResource resource, Status & status) { for (const CommandPlanGeneratedResource & existing : generated_resources_) { if (existing.source_value == resource.source_value && existing.role == resource.role) { diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h index 0025b2b0fbbf..eeb8238690ad 100644 --- a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h +++ b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h @@ -86,6 +86,18 @@ struct CommandPlanAlternateValue { std::string name; }; +struct CommandPlanActivationPublicationDiagnostic { + ValueId source_value; + ValueId alternate_value; + ValueId consumer_value; + std::string requested_format; + std::string publication_name; + std::string publication_stage; + std::string fallback_reason; + size_t alias_depth = 0; + bool mixed_fanout = false; +}; + struct CommandPlanMoeRoutingBundle { ValueId route_ids; ValueId route_weights; @@ -109,6 +121,8 @@ class CommandPlanMetadata { bool append_alternate_value(CommandPlanAlternateValue alternate, Status & status); + void append_activation_publication_diagnostic(CommandPlanActivationPublicationDiagnostic diagnostic); + bool append_moe_routing_bundle(CommandPlanMoeRoutingBundle bundle, Status & status); const CommandPlanGeneratedResource * find_generated_resource(ValueId source_value, @@ -126,11 +140,16 @@ class CommandPlanMetadata { const std::vector & alternate_values() const { return alternate_values_; } + const std::vector & activation_publication_diagnostics() const { + return activation_publication_diagnostics_; + } + const std::vector & moe_routing_bundles() const { return moe_routing_bundles_; } private: std::vector generated_resources_; std::vector alternate_values_; + std::vector activation_publication_diagnostics_; std::vector moe_routing_bundles_; }; diff --git a/ggml/src/ggml-hrx/dispatch/command-plan.h b/ggml/src/ggml-hrx/dispatch/command-plan.h index a1f1abf62168..aa6693342086 100644 --- a/ggml/src/ggml-hrx/dispatch/command-plan.h +++ b/ggml/src/ggml-hrx/dispatch/command-plan.h @@ -59,12 +59,40 @@ inline bool same_full_value_range(const Value & lhs, const Value & rhs) { return lhs.storage == rhs.storage && lhs.storage_offset == rhs.storage_offset && lhs.byte_count == rhs.byte_count; } -inline const CommandPlanAlternateValue * find_alternate_value(const Graph & graph, - const CommandPlan & plan, - ValueId graph_value, - ggml_type type, - size_t byte_count) { - const CommandPlanAlternateValue * exact = find_alternate_value(plan, graph_value, type, byte_count); +inline bool same_full_ordered_value_range(const Value & lhs, const Value & rhs) { + return same_full_value_range(lhs, rhs) && lhs.type == rhs.type && lhs.element_count == rhs.element_count && + lhs.contiguous && rhs.contiguous; +} + +inline bool same_packed_row_geometry(const Value & lhs, const Value & rhs) { + return same_full_ordered_value_range(lhs, rhs) && lhs.ne[0] > 0 && rhs.ne[0] > 0 && lhs.ne[0] == rhs.ne[0] && + lhs.element_count / lhs.ne[0] == rhs.element_count / rhs.ne[0]; +} + +inline bool same_q8_1_x4_geometry(const Value & lhs, const Value & rhs) { + return same_full_ordered_value_range(lhs, rhs) && lhs.ne[0] > 0 && rhs.ne[0] > 0 && lhs.ne[0] % 128 == 0 && + rhs.ne[0] % 128 == 0; +} + +namespace detail { + +template +inline const Resource * find_full_ordered_value_resource(const Graph & graph, + ValueId graph_value, + const Resources & resources, + SourceValue source_value, + Matches matches, + Compatible compatible) { + auto find_exact = [&](ValueId candidate_id) -> const Resource * { + for (const Resource & resource : resources) { + if (source_value(resource) == candidate_id && matches(resource)) { + return &resource; + } + } + return nullptr; + }; + + const Resource * exact = find_exact(graph_value); if (exact != nullptr) { return exact; } @@ -74,42 +102,75 @@ inline const CommandPlanAlternateValue * find_alternate_value(const Graph & return nullptr; } - auto find_if_same_range = [&](ValueId candidate_id) -> const CommandPlanAlternateValue * { + auto find_if_compatible = [&](ValueId candidate_id) -> const Resource * { const Value * candidate = graph.values().find(candidate_id); - if (candidate == nullptr || !same_full_value_range(*value, *candidate)) { + if (candidate == nullptr) { return nullptr; } - return find_alternate_value(plan, candidate_id, type, byte_count); + const Resource * resource = find_exact(candidate_id); + return resource != nullptr && compatible(*resource, *value, *candidate) ? resource : nullptr; }; - ValueId alias = value->alias_source; - for (size_t i = 0; alias.value >= 0 && i < graph.values().size(); ++i) { - const CommandPlanAlternateValue * alternate = find_if_same_range(alias); - if (alternate != nullptr) { - return alternate; - } - const Value * alias_value = graph.values().find(alias); - if (alias_value == nullptr) { - break; - } - alias = alias_value->alias_source; - } - - const CommandPlanAlternateValue * root_alternate = find_if_same_range(value->storage_root); - if (root_alternate != nullptr) { - return root_alternate; + if (!graph.has_index()) { + return nullptr; } - for (const CommandPlanAlternateValue & alternate : plan.metadata.alternate_values()) { - if (alternate.type != type || alternate.byte_count != byte_count) { - continue; + const Value * current = value; + for (size_t i = 0; i < graph.values().size(); ++i) { + const GraphNode * alias = graph.index().producer(current->id); + if (alias == nullptr || (alias->op != GGML_OP_VIEW && alias->op != GGML_OP_RESHAPE) || + alias->output != current->id || alias->inputs.size() != 1 || current->alias_source != alias->inputs[0] || + !is_layout_alias_node(graph, *alias)) { + break; } - const Value * alternate_value = graph.values().find(alternate.graph_value); - if (alternate_value != nullptr && same_full_value_range(*value, *alternate_value)) { - return &alternate; + const Value * source = graph.values().find(alias->inputs[0]); + if (source == nullptr || !same_full_ordered_value_range(*current, *source)) { + break; + } + const Resource * resource = find_if_compatible(source->id); + if (resource != nullptr) { + return resource; } + current = source; } return nullptr; } +} // namespace detail + +inline const CommandPlanAlternateValue * find_alternate_value(const Graph & graph, + const CommandPlan & plan, + ValueId graph_value, + ggml_type type, + size_t byte_count) { + return detail::find_full_ordered_value_resource( + graph, graph_value, plan.metadata.alternate_values(), + [](const CommandPlanAlternateValue & alternate) { return alternate.graph_value; }, + [=](const CommandPlanAlternateValue & alternate) { + return alternate.type == type && alternate.byte_count == byte_count; + }, + [](const CommandPlanAlternateValue & alternate, const Value & value, const Value & source) { + return alternate.type == GGML_TYPE_F16 ? same_full_ordered_value_range(value, source) : + alternate.type == GGML_TYPE_Q8_1 ? same_q8_1_x4_geometry(value, source) : + same_packed_row_geometry(value, source); + }); +} + +inline const CommandPlanGeneratedResource * find_generated_resource(const Graph & graph, + const CommandPlan & plan, + ValueId graph_value, + GeneratedResourceRole role, + size_t byte_count) { + return detail::find_full_ordered_value_resource( + graph, graph_value, plan.metadata.generated_resources(), + [](const CommandPlanGeneratedResource & resource) { return resource.source_value; }, + [=](const CommandPlanGeneratedResource & resource) { + return resource.role == role && resource.byte_count == byte_count; + }, + [](const CommandPlanGeneratedResource & resource, const Value & value, const Value & source) { + return resource.role == GeneratedResourceRole::F16K16Major ? same_packed_row_geometry(value, source) : + value.ne == source.ne; + }); +} + } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp b/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp index 24e793561970..9f43349a2fff 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp @@ -272,6 +272,24 @@ static std::string format_kernel_json(const CommandProgram & program, out << " \"byte_count\": " << program.completion_counters.byte_count << ",\n"; out << " \"count\": " << program.completion_counters.count << "\n"; out << " },\n"; + out << " \"activation_publications\": [\n"; + for (size_t i = 0; i < program.activation_publication_diagnostics.size(); ++i) { + const CommandPlanActivationPublicationDiagnostic & diagnostic = + program.activation_publication_diagnostics[i]; + if (i > 0) { + out << ",\n"; + } + out << " {\"source_value\": " << diagnostic.source_value.value + << ", \"alternate_value\": " << diagnostic.alternate_value.value + << ", \"consumer_value\": " << diagnostic.consumer_value.value + << ", \"requested_format\": \"" << json_escape(diagnostic.requested_format) + << "\", \"publication_name\": \"" << json_escape(diagnostic.publication_name) + << "\", \"publication_stage\": \"" << json_escape(diagnostic.publication_stage) + << "\", \"alias_depth\": " << diagnostic.alias_depth + << ", \"mixed_fanout\": " << (diagnostic.mixed_fanout ? "true" : "false") + << ", \"fallback_reason\": \"" << json_escape(diagnostic.fallback_reason) << "\"}"; + } + out << "\n ],\n"; out << " \"commands\": [\n"; bool wrote_any = false; if (!program.initialization_commands.empty()) { diff --git a/ggml/src/ggml-hrx/dispatch/command-program.cpp b/ggml/src/ggml-hrx/dispatch/command-program.cpp index 458930a2e8a9..4f0127c763a8 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program.cpp @@ -491,6 +491,7 @@ CommandProgram build_command_program(const Graph & graph, initialization.data, }); } + result.activation_publication_diagnostics = plan.metadata.activation_publication_diagnostics(); return result; } diff --git a/ggml/src/ggml-hrx/dispatch/command-program.h b/ggml/src/ggml-hrx/dispatch/command-program.h index 42796c30aa86..3b7cd8fd7465 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program.h +++ b/ggml/src/ggml-hrx/dispatch/command-program.h @@ -78,6 +78,7 @@ struct CommandProgram { TransientPlan transients; CompletionCounterPlan completion_counters; std::vector constant_initializations; + std::vector activation_publication_diagnostics; Status status; bool valid() const { return status.success(); } diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp index 3c668d0c9cf5..21990fe038fa 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp @@ -35,7 +35,7 @@ static bool try_match_registration(const Graph & graph, DispatchMatch & match, DispatchMatchDiagnostics * diagnostics) { const DispatchMatchContext context = { - graph, node, node_index, covered_nodes, plan, next_plan_value, + graph, node, node_index, covered_nodes, plan, next_plan_value, ®istry, }; return registry.match(context, match, diagnostics); } diff --git a/ggml/src/ggml-hrx/dispatch/dispatch.h b/ggml/src/ggml-hrx/dispatch/dispatch.h index 0e5ea1b0cfa9..a59879faef9c 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch.h +++ b/ggml/src/ggml-hrx/dispatch/dispatch.h @@ -27,8 +27,14 @@ inline constexpr const char kQ5KSymmetricI8K256Row64Layout[] = "q5k-symi8-k25 inline constexpr const char kQ6KI8K32Row64Layout[] = "q6k-i8-k32-row64"; inline constexpr const char kQ6KPackedK256Row64ScaleRowLayout[] = "q6k-packed-k256-row64-scalerow"; +enum class WorkloadSpecialization : uint8_t { + Exact, + Dynamic, +}; + struct KernelSpecialization { uint64_t kernel_id = kUncatalogedKernelId; + WorkloadSpecialization workload_specialization = WorkloadSpecialization::Exact; std::map integer_parameters; std::map compile_parameters; }; diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.cpp new file mode 100644 index 000000000000..df6b24972b37 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.cpp @@ -0,0 +1,502 @@ +#include "dispatch-activation-publication.h" + +#include "dispatch/command-plan.h" +#include "ggml.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static bool is_ordered_alias_node(const Graph & graph, const GraphNode & node, const Value & input) { + if ((node.op != GGML_OP_VIEW && node.op != GGML_OP_RESHAPE) || node.inputs.size() != 1 || + node.inputs.front() != input.id || !is_layout_alias_node(graph, node)) { + return false; + } + const Value * output = graph.values().find(node.output); + return output != nullptr && output->alias_source == input.id && same_full_ordered_value_range(input, *output); +} + +static bool publication_values_compatible(CommonActivationPublicationFormat format, + const Value & source, + const Value & candidate) { + if (format == CommonActivationPublicationFormat::F16Row) { + return same_full_ordered_value_range(source, candidate); + } + if (format == CommonActivationPublicationFormat::Q8_1X4) { + return same_q8_1_x4_geometry(source, candidate); + } + return same_packed_row_geometry(source, candidate); +} + +static uint32_t publication_format_mask(CommonActivationPublicationFormat format) { + switch (format) { + case CommonActivationPublicationFormat::F16Row: + return DispatchActivationInputF16Row; + case CommonActivationPublicationFormat::F16K16Major: + return DispatchActivationInputF16K16Major; + case CommonActivationPublicationFormat::Q8_1X4: + return DispatchActivationInputQ8_1X4; + case CommonActivationPublicationFormat::SymmetricI4K32: + return DispatchActivationInputSymmetricI4K32; + case CommonActivationPublicationFormat::None: + return DispatchActivationInputNone; + } + return DispatchActivationInputNone; +} + +static CommonActivationPublicationFormat publication_format_from_mask(uint32_t format) { + switch (format) { + case DispatchActivationInputF16Row: + return CommonActivationPublicationFormat::F16Row; + case DispatchActivationInputF16K16Major: + return CommonActivationPublicationFormat::F16K16Major; + case DispatchActivationInputQ8_1X4: + return CommonActivationPublicationFormat::Q8_1X4; + case DispatchActivationInputSymmetricI4K32: + return CommonActivationPublicationFormat::SymmetricI4K32; + default: + return CommonActivationPublicationFormat::None; + } +} + +static int publication_use_priority(DispatchActivationConsumerUse use) { + switch (use) { + case DispatchActivationConsumerUse::BandwidthLimited: + return 400; + case DispatchActivationConsumerUse::Tiled: + return 300; + case DispatchActivationConsumerUse::LowRow: + return 250; + case DispatchActivationConsumerUse::Row: + return 200; + } + return 0; +} + +static int publication_format_priority(uint32_t format) { + switch (format) { + case DispatchActivationInputQ8_1X4: + return 40; + case DispatchActivationInputF16K16Major: + return 30; + case DispatchActivationInputSymmetricI4K32: + return 20; + case DispatchActivationInputF16Row: + return 10; + default: + return 0; + } +} + +static uint32_t publication_consumer_source_mask(DispatchSource source) { + switch (source) { + case DispatchSource::Common: + return DispatchActivationConsumerSourceCommon; + case DispatchSource::Llm: + return DispatchActivationConsumerSourceLlm; + case DispatchSource::Qwen: + return DispatchActivationConsumerSourceQwen; + } + return 0; +} + +static bool publication_has_mixed_fanout(const Graph & graph, + const Value & produced_value, + const Value & consumer_value, + const GraphNode & selected_consumer) { + const Value * current = &consumer_value; + const GraphNode * allowed_consumer = &selected_consumer; + for (size_t depth = 0; depth <= graph.values().size(); ++depth) { + for (const GraphNode * consumer : graph.index().consumers(current->id)) { + if (consumer != nullptr && consumer != allowed_consumer) { + return true; + } + } + if (current->id == produced_value.id) { + return false; + } + const GraphNode * alias = graph.index().producer(current->id); + if (alias == nullptr || alias->inputs.size() != 1) { + return true; + } + const Value * source = graph.values().find(alias->inputs.front()); + if (source == nullptr) { + return true; + } + allowed_consumer = alias; + current = source; + } + return true; +} + +static size_t align_up(size_t value, size_t alignment) { + return value > std::numeric_limits::max() - (alignment - 1) ? + 0 : + (value + alignment - 1) / alignment * alignment; +} + +static size_t add_and_align(size_t value, size_t increment, size_t alignment) { + return value > std::numeric_limits::max() - increment ? 0 : align_up(value + increment, alignment); +} + +static bool checked_element_count(const Value & value, size_t & element_count) { + if (value.type != GGML_TYPE_F32 || !value.contiguous || value.element_count <= 0 || value.ne[0] <= 0 || + value.element_count % value.ne[0] != 0) { + return false; + } + element_count = static_cast(value.element_count); + return static_cast(element_count) == value.element_count; +} + +static bool set_publication_layout(CommonActivationPublicationFormat format, + const Value & value, + CommonActivationPublication & publication) { + size_t element_count = 0; + if (!checked_element_count(value, element_count)) { + return false; + } + + switch (format) { + case CommonActivationPublicationFormat::F16Row: + case CommonActivationPublicationFormat::F16K16Major: + if (element_count > std::numeric_limits::max() / sizeof(ggml_fp16_t)) { + return false; + } + publication.byte_count = element_count * sizeof(ggml_fp16_t); + publication.payload_bytes = publication.byte_count; + return true; + case CommonActivationPublicationFormat::Q8_1X4: + { + if (value.ne[0] % 128 != 0) { + return false; + } + const size_t row_count = element_count / static_cast(value.ne[0]); + const size_t row_bytes = ggml_row_size(GGML_TYPE_Q8_1, value.ne[0]); + if (row_count > std::numeric_limits::max() / row_bytes) { + return false; + } + publication.byte_count = row_count * row_bytes; + publication.payload_bytes = publication.byte_count; + return publication.byte_count > 0; + } + case CommonActivationPublicationFormat::SymmetricI4K32: + { + if (value.ne[0] % 64 != 0 || element_count % 8 != 0) { + return false; + } + const size_t payload_bytes = element_count / 2; + const size_t metadata_bytes = element_count / 8; + const size_t scales_offset = align_up(payload_bytes, 256); + const size_t sums_offset = scales_offset == 0 ? 0 : add_and_align(scales_offset, metadata_bytes, 256); + const size_t total_bytes = sums_offset == 0 ? 0 : add_and_align(sums_offset, metadata_bytes, 256); + if (total_bytes == 0) { + return false; + } + publication.byte_count = total_bytes; + publication.payload_bytes = payload_bytes; + publication.scales_offset = scales_offset; + publication.scales_bytes = metadata_bytes; + publication.sums_offset = sums_offset; + publication.sums_bytes = metadata_bytes; + return true; + } + case CommonActivationPublicationFormat::None: + return false; + } + return false; +} + +static bool reserve_activation_output(const DispatchMatchContext & context, + DispatchMatch & match, + const Value & layout_value, + const char * transient_name, + CommonActivationPublication & publication) { + if (transient_name == nullptr || !set_publication_layout(publication.format, layout_value, publication)) { + publication = {}; + return false; + } + + const size_t value_offset = match.transients.size() + match.completion_counter_requests.size(); + if (context.next_plan_value.value < 0 || value_offset > static_cast(std::numeric_limits::max()) || + context.next_plan_value.value > std::numeric_limits::max() - static_cast(value_offset)) { + publication = {}; + return false; + } + publication.alternate_value = ValueId(context.next_plan_value.value + static_cast(value_offset)); + match.transients.push_back({ publication.alternate_value, transient_name, publication.byte_count, 256 }); + return true; +} + +} // namespace + +const char * common_activation_publication_format_name(CommonActivationPublicationFormat format) { + switch (format) { + case CommonActivationPublicationFormat::F16Row: + return "f16-row"; + case CommonActivationPublicationFormat::F16K16Major: + return "f16-k16-major"; + case CommonActivationPublicationFormat::Q8_1X4: + return "q8-1-x4"; + case CommonActivationPublicationFormat::SymmetricI4K32: + return "symmetric-i4-k32"; + case CommonActivationPublicationFormat::None: + return "none"; + } + return "none"; +} + +const char * common_activation_publication_fallback_reason_name(CommonActivationPublicationFallbackReason reason) { + switch (reason) { + case CommonActivationPublicationFallbackReason::None: + return "none"; + case CommonActivationPublicationFallbackReason::MissingGraphIndex: + return "missing-graph-index"; + case CommonActivationPublicationFallbackReason::NoEnabledCandidates: + return "no-enabled-candidates"; + case CommonActivationPublicationFallbackReason::NoQualifiedConsumer: + return "no-qualified-consumer"; + case CommonActivationPublicationFallbackReason::IncompatibleAlias: + return "incompatible-alias"; + } + return "unknown"; +} + +CommonActivationPublicationPlan common_select_activation_publication_plan( + const DispatchMatchContext & context, + const Value & produced_value, + const CommonActivationPublicationCapabilities & capabilities) { + CommonActivationPublicationPlan result; + if (!context.graph.has_index()) { + result.fallback_reason = CommonActivationPublicationFallbackReason::MissingGraphIndex; + return result; + } + if (capabilities.formats == DispatchActivationInputNone) { + result.fallback_reason = CommonActivationPublicationFallbackReason::NoEnabledCandidates; + return result; + } + if (context.registry == nullptr) { + result.fallback_reason = CommonActivationPublicationFallbackReason::NoQualifiedConsumer; + return result; + } + + struct TraversedValue { + const Value * value = nullptr; + size_t depth = 0; + uint32_t eligible_formats = DispatchActivationInputNone; + }; + struct Selection { + CommonActivationPublicationDemand demand; + int priority = 0; + }; + std::vector values = { { &produced_value, 0, capabilities.formats } }; + std::array best_by_format; + bool incompatible_alias = false; + for (size_t value_index = 0; value_index < values.size(); ++value_index) { + const Value & value = *values[value_index].value; + for (const GraphNode * consumer : context.graph.index().consumers(value.id)) { + if (consumer == nullptr) { + continue; + } + const auto & classifiers = context.registry->activation_consumers_for_root(consumer->op); + for (const DispatchActivationConsumerRegistration & classifier : classifiers) { + if ((capabilities.consumer_sources & publication_consumer_source_mask(classifier.source)) == 0 || + (classifier.traits & capabilities.required_consumer_traits) != + capabilities.required_consumer_traits || + (classifier.producer_requirements & capabilities.producer_requirements) != + classifier.producer_requirements || + !classifier.accepts(context, *consumer, value)) { + continue; + } + uint32_t accepted = classifier.accepted_formats & values[value_index].eligible_formats; + while (accepted != 0) { + const uint32_t format = accepted & (~accepted + 1); + accepted &= ~format; + const int priority = publication_use_priority(classifier.use) * 100 + + publication_format_priority(format); + const size_t format_index = static_cast(publication_format_from_mask(format)) - 1; + Selection & best = best_by_format[format_index]; + if (priority > best.priority) { + best.priority = priority; + best.demand = { publication_format_from_mask(format), + &value, + consumer, + values[value_index].depth, + publication_has_mixed_fanout(context.graph, produced_value, value, *consumer), + CommonActivationPublicationFallbackReason::None }; + } + } + } + + if ((values[value_index].eligible_formats & capabilities.alias_traversal_formats) == 0 || + values[value_index].depth >= context.graph.values().size() || + !is_ordered_alias_node(context.graph, *consumer, value)) { + continue; + } + const Value * alias = context.graph.values().find(consumer->output); + if (alias == nullptr) { + continue; + } + uint32_t alias_formats = DispatchActivationInputNone; + for (uint32_t format = DispatchActivationInputF16Row; + format <= DispatchActivationInputSymmetricI4K32; format <<= 1) { + if ((values[value_index].eligible_formats & capabilities.alias_traversal_formats & format) != 0 && + publication_values_compatible(publication_format_from_mask(format), produced_value, *alias)) { + alias_formats |= format; + } + } + if (alias_formats == DispatchActivationInputNone) { + incompatible_alias = true; + continue; + } + TraversedValue * existing = nullptr; + for (TraversedValue & traversed : values) { + if (traversed.value->id == alias->id) { + existing = &traversed; + break; + } + } + if (existing == nullptr) { + values.push_back({ alias, values[value_index].depth + 1, alias_formats }); + } else { + existing->eligible_formats |= alias_formats; + } + } + } + std::vector selections; + for (const Selection & selection : best_by_format) { + if (selection.priority != 0) { + selections.push_back(selection); + } + } + std::stable_sort(selections.begin(), selections.end(), [](const Selection & lhs, const Selection & rhs) { + return lhs.priority > rhs.priority; + }); + if (selections.empty()) { + result.fallback_reason = incompatible_alias ? CommonActivationPublicationFallbackReason::IncompatibleAlias : + CommonActivationPublicationFallbackReason::NoQualifiedConsumer; + return result; + } + + result.publications[result.publication_count++] = selections.front().demand; + const CommonActivationPublicationDemand & preferred = selections.front().demand; + const uint32_t preferred_format = publication_format_mask(preferred.format); + if ((capabilities.co_publication_formats & preferred_format) != 0) { + static constexpr uint32_t kFormatPriority[] = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputF16K16Major, + DispatchActivationInputSymmetricI4K32, + DispatchActivationInputF16Row, + }; + for (uint32_t format : kFormatPriority) { + if (format == preferred_format || (capabilities.co_publication_formats & format) == 0 || + result.publication_count >= result.publications.size()) { + continue; + } + const CommonActivationPublicationFormat publication_format = publication_format_from_mask(format); + const auto selection = std::find_if( + selections.begin(), selections.end(), [publication_format](const Selection & candidate) { + return candidate.demand.format == publication_format; + }); + if (selection != selections.end()) { + result.publications[result.publication_count++] = selection->demand; + } else if (publication_values_compatible(publication_format, produced_value, *preferred.consumer_value)) { + CommonActivationPublicationDemand companion = preferred; + companion.format = publication_format; + result.publications[result.publication_count++] = companion; + } + } + } + return result; +} + +bool common_reserve_activation_publication(const DispatchMatchContext & context, + DispatchMatch & match, + const Value & produced_value, + const CommonActivationPublicationDemand & demand, + const char * transient_name, + CommonActivationPublication & publication) { + publication = {}; + if (!demand.matched() || transient_name == nullptr || + !publication_values_compatible(demand.format, produced_value, *demand.consumer_value)) { + return false; + } + + publication.format = demand.format; + publication.source_value = produced_value.id; + publication.consumer_value = demand.consumer_value->id; + publication.alias_depth = demand.alias_depth; + publication.mixed_fanout = demand.mixed_fanout; + return reserve_activation_output(context, match, *demand.consumer_value, transient_name, publication); +} + +bool common_reserve_private_activation_output(const DispatchMatchContext & context, + DispatchMatch & match, + const Value & produced_value, + CommonActivationPublicationFormat format, + const char * transient_name, + CommonActivationPublication & publication) { + publication = {}; + publication.format = format; + publication.source_value = produced_value.id; + publication.consumer_value = produced_value.id; + return reserve_activation_output(context, match, produced_value, transient_name, publication); +} + +bool common_append_activation_publication(DispatchMatch & match, + const CommonActivationPublication & publication, + const char * metadata_name) { + if (!publication.matched() || metadata_name == nullptr) { + return false; + } + + Status status; + bool recorded = false; + switch (publication.format) { + case CommonActivationPublicationFormat::F16Row: + recorded = match.metadata.append_alternate_value({ publication.source_value, publication.alternate_value, + GGML_TYPE_F16, publication.byte_count, metadata_name }, + status); + break; + case CommonActivationPublicationFormat::F16K16Major: + recorded = match.metadata.append_generated_resource({ publication.source_value, + GeneratedResourceRole::F16K16Major, + publication.alternate_value, + publication.byte_count, + {} }, + status); + break; + case CommonActivationPublicationFormat::Q8_1X4: + recorded = match.metadata.append_alternate_value({ publication.source_value, publication.alternate_value, + GGML_TYPE_Q8_1, publication.byte_count, metadata_name }, + status); + break; + case CommonActivationPublicationFormat::SymmetricI4K32: + recorded = match.metadata.append_alternate_value({ publication.source_value, publication.alternate_value, + GGML_TYPE_COUNT, publication.byte_count, metadata_name }, + status); + break; + case CommonActivationPublicationFormat::None: + break; + } + if (!recorded) { + match.status.append(status); + } else { + match.metadata.append_activation_publication_diagnostic({ publication.source_value, + publication.alternate_value, + publication.consumer_value, + common_activation_publication_format_name( + publication.format), + metadata_name, + "published", + "none", + publication.alias_depth, + publication.mixed_fanout }); + } + return recorded; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.h new file mode 100644 index 000000000000..5d9646b06982 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-activation-publication.h @@ -0,0 +1,125 @@ +#pragma once + +#include "../dispatch-registry.h" + +#include +#include + +namespace ggml::hrx { + +enum class CommonActivationPublicationFormat { + None, + F16Row, + F16K16Major, + Q8_1X4, + SymmetricI4K32, +}; + +enum class CommonActivationConsumerTraversal { + Direct, + FullRangeOrderedAliases, +}; + +enum class CommonActivationPublicationFallbackReason { + None, + MissingGraphIndex, + NoEnabledCandidates, + NoQualifiedConsumer, + IncompatibleAlias, +}; + +struct CommonActivationPublicationCapabilities { + uint32_t formats = DispatchActivationInputNone; + uint32_t alias_traversal_formats = DispatchActivationInputNone; + uint32_t producer_requirements = DispatchActivationProducerRequirementNone; + uint32_t co_publication_formats = DispatchActivationInputNone; + uint32_t consumer_sources = DispatchActivationConsumerSourceAll; + uint32_t required_consumer_traits = DispatchActivationConsumerTraitNone; +}; + +struct CommonActivationPublicationDemand { + CommonActivationPublicationFormat format = CommonActivationPublicationFormat::None; + const Value * consumer_value = nullptr; + const GraphNode * consumer = nullptr; + size_t alias_depth = 0; + bool mixed_fanout = false; + CommonActivationPublicationFallbackReason fallback_reason = CommonActivationPublicationFallbackReason::None; + + bool matched() const { + return format != CommonActivationPublicationFormat::None && consumer_value != nullptr && consumer != nullptr; + } +}; + +struct CommonActivationPublicationPlan { + std::array publications; + size_t publication_count = 0; + CommonActivationPublicationFallbackReason fallback_reason = + CommonActivationPublicationFallbackReason::None; + + bool matched() const { return publication_count != 0; } + + const CommonActivationPublicationDemand & preferred() const { return publications.front(); } + + const CommonActivationPublicationDemand * find(CommonActivationPublicationFormat format) const { + for (size_t i = 0; i < publication_count; ++i) { + if (publications[i].format == format) { + return &publications[i]; + } + } + return nullptr; + } +}; + +struct CommonActivationPublication { + CommonActivationPublicationFormat format = CommonActivationPublicationFormat::None; + ValueId source_value; + ValueId alternate_value; + size_t byte_count = 0; + size_t payload_bytes = 0; + size_t scales_offset = 0; + size_t scales_bytes = 0; + size_t sums_offset = 0; + size_t sums_bytes = 0; + ValueId consumer_value; + size_t alias_depth = 0; + bool mixed_fanout = false; + + bool matched() const { + return format != CommonActivationPublicationFormat::None && source_value.value >= 0 && + alternate_value.value >= 0 && byte_count > 0; + } + + DispatchBinding binding() const { return { alternate_value, 0, byte_count }; } + + DispatchBinding binding(size_t offset, size_t length) const { return { alternate_value, offset, length }; } +}; + +const char * common_activation_publication_format_name(CommonActivationPublicationFormat format); + +const char * common_activation_publication_fallback_reason_name(CommonActivationPublicationFallbackReason reason); + +CommonActivationPublicationPlan common_select_activation_publication_plan( + const DispatchMatchContext & context, + const Value & produced_value, + const CommonActivationPublicationCapabilities & capabilities); + +bool common_reserve_activation_publication(const DispatchMatchContext & context, + DispatchMatch & match, + const Value & produced_value, + const CommonActivationPublicationDemand & demand, + const char * transient_name, + CommonActivationPublication & publication); + +bool common_reserve_private_activation_output(const DispatchMatchContext & context, + DispatchMatch & match, + const Value & produced_value, + CommonActivationPublicationFormat format, + const char * transient_name, + CommonActivationPublication & publication); + +// A metadata conflict is appended to match.status and invalidates the entire match; callers must not fall back. +bool common_append_activation_publication(DispatchMatch & match, + const CommonActivationPublication & publication, + const char * metadata_name); + +} // namespace ggml::hrx 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 84f34aa28c79..c42ac59a441a 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp @@ -1,5 +1,6 @@ #include "dispatch-binary.h" +#include "dispatch-activation-publication.h" #include "dispatch-layout-utils.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" @@ -15,6 +16,10 @@ namespace { static constexpr KernelCatalogRef kBinaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_f32"); static constexpr KernelCatalogRef kBinaryBcF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_bc_f32"); +static constexpr KernelCatalogRef kBinaryF32PublishQ8_1X4Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_f32_publish_q8_1_x4"); +static constexpr KernelCatalogRef kBinaryBcF32PublishQ8_1X4Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_bc_f32_publish_q8_1_x4"); static constexpr KernelCatalogRef kBinarySwiGluSymmetricI4K32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_swiglu_symmetric_i4_k32"); @@ -128,6 +133,12 @@ static bool binary_kind_allows_broadcast(BinaryKind kind, const Value & lhs, con return false; } +static bool binary_q8_publication_enabled(BinaryKind kind, const Value & output) { + return (kind == BinaryKind::Add || kind == BinaryKind::Mul) && output.ne[0] >= 128 && output.ne[0] <= 32768 && + output.ne[0] % 128 == 0 && output.ne[1] >= 1 && output.ne[1] <= 5 && output.ne[2] == 1 && + output.ne[3] == 1 && output.element_count == output.ne[0] * output.ne[1]; +} + static uint32_t broadcast_dim_flag(const Value & source, const Value & output, int dim) { return source.ne[dim] == 1 && output.ne[dim] != 1 ? 1 : 0; } @@ -216,8 +227,7 @@ static bool match_binary_swiglu_symmetric_i4_dispatch(const DispatchMatchContext !packed_f32_layout(*output) || output->alias_source.value >= 0 || !supported_source_layout(context.graph, *lhs) || !supported_source_layout(context.graph, *rhs) || !binary_noalias_storage_is_safe(*lhs, *rhs, *output) || output->ne[0] < 256 || output->ne[0] > 32768 || - output->ne[0] % 64 != 0 || output->element_count <= 0 || output->element_count % output->ne[0] != 0 || - !common_has_symmetric_i4_lowrow_consumer(context.graph, *output)) { + output->ne[0] % 64 != 0 || output->element_count <= 0 || output->element_count % output->ne[0] != 0) { return false; } @@ -226,12 +236,20 @@ static bool match_binary_swiglu_symmetric_i4_dispatch(const DispatchMatchContext if (token_count < 1 || token_count > 16) { return false; } - const CommonSymmetricI4ActivationLayout activation_layout = - common_symmetric_i4_activation_layout(input_size, token_count); - if (activation_layout.total_bytes == 0) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputSymmetricI4K32, + DispatchActivationInputSymmetricI4K32, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *output, capabilities); + const CommonActivationPublicationDemand demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *output, demand, + kCommonSymmetricI4K32ActivationAlternateName, publication) || + !common_append_activation_publication(match, publication, kCommonSymmetricI4K32ActivationAlternateName)) { return false; } - const ValueId activation = context.next_plan_value; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kBinarySwiGluSymmetricI4K32Kernel); @@ -240,22 +258,12 @@ static bool match_binary_swiglu_symmetric_i4_dispatch(const DispatchMatchContext dispatch.kernel.compile_parameters.emplace("ggml.binary_swiglu_symmetric_i4.token_count", std::to_string(token_count)); bind_binary_buffers(dispatch, *lhs, lhs->byte_count, *rhs, rhs->byte_count, *output); - dispatch.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); - - Status metadata_status; - if (!match.metadata.append_alternate_value({ output->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, - kCommonSymmetricI4K32ActivationAlternateName }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(publication.binding(0, publication.payload_bytes)); + dispatch.bindings.push_back(publication.binding(publication.scales_offset, publication.scales_bytes)); + dispatch.bindings.push_back(publication.binding(publication.sums_offset, publication.sums_bytes)); match.covered_nodes.push_back(context.root_index); match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back( - { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); return match.status.success(); } @@ -289,10 +297,26 @@ static bool match_binary_f32_dispatch(const DispatchMatchContext & context, Disp return false; } + const bool same_shape_inputs = same_shape(*lhs, *output) && same_shape(*rhs, *output); + const CommonActivationPublicationCapabilities capabilities = { + binary_q8_publication_enabled(params->op, *output) ? DispatchActivationInputQ8_1X4 : + DispatchActivationInputNone, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *output, capabilities); + const CommonActivationPublicationDemand * q8_demand = + publication_plan.find(CommonActivationPublicationFormat::Q8_1X4); + const bool publish_q8 = q8_demand != nullptr; + Dispatch dispatch; - if (same_shape(*lhs, *output) && same_shape(*rhs, *output)) { - dispatch.kernel = make_kernel_specialization(kBinaryF32Kernel); + if (same_shape_inputs) { + 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]); + } dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.op", std::to_string(binary_kind_config_value(params->op))); } else { @@ -302,18 +326,34 @@ static bool match_binary_f32_dispatch(const DispatchMatchContext & context, Disp if (!binary_kind_allows_broadcast(params->op, *lhs, *rhs, *output)) { return false; } - dispatch.kernel = make_kernel_specialization(kBinaryBcF32Kernel); - add_binary_shape_parameters(dispatch, *lhs, *rhs, *output); + dispatch.kernel = make_kernel_specialization(publish_q8 ? kBinaryBcF32PublishQ8_1X4Kernel : kBinaryBcF32Kernel); + if (publish_q8) { + dispatch.kernel.integer_parameters.emplace("token_count", output->ne[1]); + dispatch.kernel.integer_parameters.emplace("hidden_size", output->ne[0]); + dispatch.kernel.integer_parameters.emplace("src0_element_count", lhs->element_count); + dispatch.kernel.integer_parameters.emplace("src1_element_count", rhs->element_count); + } else { + add_binary_shape_parameters(dispatch, *lhs, *rhs, *output); + } dispatch.kernel.compile_parameters.emplace("ggml.binary_bc_f32.op", std::to_string(binary_kind_config_value(params->op))); add_broadcast_config(dispatch, "ggml.binary_bc_f32.", "src0", *lhs, *output); add_broadcast_config(dispatch, "ggml.binary_bc_f32.", "src1", *rhs, *output); } bind_binary_buffers(dispatch, *lhs, lhs_byte_count, *rhs, rhs_byte_count, *output); + if (publish_q8) { + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *output, *q8_demand, "common.binary.q8_1_x4", + publication) || + !common_append_activation_publication(match, publication, "common.binary.q8_1_x4")) { + return false; + } + dispatch.bindings.push_back(publication.binding()); + } match.covered_nodes.push_back(context.root_index); match.dispatches.push_back(std::move(dispatch)); - return true; + return match.status.success(); } static void register_binary_dispatch_for(DispatchRegistryBuilder & registry, ggml_op root_op) { 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 cb531ccd1cfe..5f091f40e564 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp @@ -72,6 +72,7 @@ static bool make_copy_f32_dispatch(const Value & source, const Value & output, D size_t source_span = source.byte_count; if (packed_f32_layout(source)) { dispatch.kernel = make_kernel_specialization(kCopyF32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; } else { if (!strided_f32_storage_span_bytes(source, source_span)) { return false; @@ -174,9 +175,12 @@ static bool match_concat_dim0_f32_dispatch(const DispatchMatchContext & context, Dispatch dispatch; if (packed_f32_layout(*lhs) && packed_f32_layout(*rhs)) { dispatch.kernel = make_kernel_specialization(kConcatDim0F32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + dispatch.kernel.integer_parameters.emplace("row_count", row_count); + dispatch.kernel.integer_parameters.emplace("lhs_span", lhs->element_count); + dispatch.kernel.integer_parameters.emplace("rhs_span", rhs->element_count); dispatch.kernel.compile_parameters.emplace("ggml.concat_dim0_f32.lhs_width", std::to_string(lhs->ne[0])); dispatch.kernel.compile_parameters.emplace("ggml.concat_dim0_f32.rhs_width", std::to_string(rhs->ne[0])); - dispatch.kernel.compile_parameters.emplace("ggml.concat_dim0_f32.row_count", std::to_string(row_count)); } else { const size_t lhs_span = lhs->byte_count / sizeof(float); const size_t rhs_span = rhs->byte_count / sizeof(float); 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 585cb045f8fe..c1fec1a8a8ce 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 @@ -1,5 +1,6 @@ #include "dispatch-flash-attention.h" +#include "dispatch-activation-publication.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -39,6 +40,11 @@ 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; +} + static bool is_supported_token_count(int64_t token_count) { return token_count >= 1 && token_count <= 2048; } @@ -114,19 +120,18 @@ static std::string to_config_value(int64_t value) { return std::to_string(value); } -static size_t q8_1_x4_byte_count(int64_t row_count, int64_t hidden_size) { - if (row_count <= 0 || hidden_size <= 0) { - return 0; - } - // Packed Q8 stores four 32-element blocks in each physical group. - const int64_t padded_hidden_size = (hidden_size + 127) / 128 * 128; - return static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, padded_hidden_size); -} - static int64_t ceil_div(int64_t value, int64_t divisor) { return (value + divisor - 1) / divisor; } +static int64_t decode_key_value_capacity(int64_t token_count) { + int64_t capacity = kDecodeKvTileSize; + while (capacity < token_count) { + capacity *= 2; + } + return capacity; +} + static ValueId match_value(const DispatchMatchContext & context, const DispatchMatch & dispatch_match, int32_t offset) { return ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size()) + static_cast(dispatch_match.completion_counter_requests.size()) + offset); @@ -337,7 +342,7 @@ static DecodeSplitFlashAttentionMatch match_decode_split_flash_attention_f32_f16 match.output_layout = find_single_layout_alias_consumer(graph, output->id); match.query_token_count = query_token_count; match.key_value_token_count = key_value_token_count; - match.key_value_capacity = ceil_div(key_value_token_count, kDecodeKvTileSize) * kDecodeKvTileSize; + match.key_value_capacity = decode_key_value_capacity(key_value_token_count); match.query_head_count = query_head_count; match.key_value_head_count = key_value_head_count; match.qk_head_size = qk_head_size; @@ -409,8 +414,9 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext Dispatch copy; copy.kernel = make_kernel_specialization(kCopyTransposeF16Kernel); - copy.kernel.compile_parameters.emplace("ggml.copy_transpose_f16.row_count", to_config_value(match.key_value_token_count)); - copy.kernel.compile_parameters.emplace("ggml.copy_transpose_f16.column_count", to_config_value(columns)); + copy.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + copy.kernel.integer_parameters.emplace("row_count", match.key_value_token_count); + copy.kernel.integer_parameters.emplace("column_count", columns); copy.bindings.push_back({ match.value->id, 0, bytes }); copy.bindings.push_back({ transposed, 0, bytes }); dispatch_match.dispatches.push_back(std::move(copy)); @@ -418,35 +424,6 @@ static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext return { transposed, 0, bytes }; } -static bool has_qualified_f16_consumer(const Graph & graph, const FlashAttentionMatch & match) { - if (!graph.has_index() || match.query_token_count < 128) { - return false; - } - const Value * published = match.output_layout != nullptr ? graph_value(graph, match.output_layout->output) : - match.output; - if (published == nullptr || published->type != GGML_TYPE_F32 || !published->contiguous) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(published->id)) { - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(graph, consumer, kFlashAttentionF32F16WmmaKernel, false); - if (!projection.matched() || projection.input->id != published->id || - projection.input_size != published->ne[0] || projection.token_count != match.query_token_count || - projection.input_size % 256 != 0 || projection.output_size % 64 != 0 || - projection.weight->alias_source.value >= 0) { - continue; - } - if (projection.weight->type == GGML_TYPE_Q6_K && projection.token_count % 128 == 0) { - return true; - } - if (projection.weight->type == GGML_TYPE_Q4_K && projection.token_count % 256 == 0 && - projection.output_size >= projection.input_size / 4) { - return true; - } - } - return false; -} - 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) { @@ -514,6 +491,9 @@ static bool match_flash_attention_gate_dispatch(const DispatchMatchContext & con Dispatch dispatch; dispatch.kernel = make_kernel_specialization(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("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, @@ -537,10 +517,21 @@ static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & return false; } - const bool publish_f16 = has_qualified_f16_consumer(context.graph, match); + const CommonActivationPublicationCapabilities capabilities = { + 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); + 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("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, @@ -554,18 +545,13 @@ static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_f16) { - const size_t bytes = match.output->byte_count / 2; - const ValueId f16_output(context.next_plan_value.value + - static_cast(dispatch_match.transients.size())); constexpr const char * name = "common.flash_attention.f16"; - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { match.output->id, f16_output, GGML_TYPE_F16, bytes, name }, status)) { - dispatch_match.status.append(status); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, *demand, name, publication) || + !common_append_activation_publication(dispatch_match, publication, name)) { return false; } - dispatch.bindings.push_back({ f16_output, 0, bytes }); - dispatch_match.transients.push_back({ f16_output, name, bytes, 256 }); + dispatch.bindings.push_back(publication.binding()); } dispatch_match.covered_nodes.push_back(context.root_index); @@ -587,6 +573,22 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa return false; } + const Value * publication_subject = + match.output_layout != nullptr ? graph_value(context.graph, match.output_layout->output) : nullptr; + if (match.output_layout != nullptr && + (match.output_layout->op != GGML_OP_RESHAPE || publication_subject == nullptr || + !same_full_ordered_value_range(*match.output, *publication_subject))) { + return false; + } + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = publication_subject != nullptr ? + common_select_activation_publication_plan(context, *publication_subject, capabilities) : + CommonActivationPublicationPlan{}; + const CommonActivationPublicationDemand * demand = + publication_plan.find(CommonActivationPublicationFormat::Q8_1X4); + const int64_t key_value_block_count = ceil_div(match.key_value_capacity, kDecodeKvTileSize); const size_t partial_scalar_count = static_cast(match.key_value_head_count) * static_cast(key_value_block_count) * @@ -596,9 +598,7 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa const size_t partial_output_bytes = partial_value_count * sizeof(ggml_fp16_t); const int64_t query_hidden_size = match.query_head_count * match.qk_head_size; const int64_t output_hidden_size = match.query_head_count * match.value_head_size; - const size_t q8_row_bytes = q8_1_x4_byte_count(1, output_hidden_size); - const size_t q8_output_bytes = q8_1_x4_byte_count(match.query_token_count, output_hidden_size); - if (partial_scalar_bytes == 0 || partial_output_bytes == 0 || q8_row_bytes == 0 || q8_output_bytes == 0) { + if (partial_scalar_bytes == 0 || partial_output_bytes == 0) { return false; } @@ -606,7 +606,6 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa const ValueId partial_sum = match_value(context, dispatch_match, 1); const ValueId partial_output = match_value(context, dispatch_match, 2); const ValueId completion_counter = match_value(context, dispatch_match, 3); - const ValueId q8_output = match_value(context, dispatch_match, 4); dispatch_match.transients.push_back( { partial_max, "common.decode.flash_attention.partial_max", partial_scalar_bytes, 256 }); @@ -614,21 +613,29 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa { partial_sum, "common.decode.flash_attention.partial_sum", partial_scalar_bytes, 256 }); dispatch_match.transients.push_back( { partial_output, "common.decode.flash_attention.partial_output", partial_output_bytes, 256 }); - dispatch_match.transients.push_back( - { q8_output, "common.decode.flash_attention.next_q8_output", q8_output_bytes, 256 }); dispatch_match.completion_counter_requests.push_back({ completion_counter, "common.decode.flash_attention.completion_counter", static_cast(match.key_value_head_count), }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value({ match.output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, - "common.decode.flash_attention.next_q8_output" }, - metadata_status)) { - dispatch_match.status.append(metadata_status); + CommonActivationPublication publication; + constexpr const char * publication_name = "common.decode.flash_attention.next_q8_output"; + Value private_layout = publication_subject != nullptr ? *publication_subject : *match.output; + if (publication_subject == nullptr) { + private_layout.ne = { output_hidden_size * match.query_token_count, 1, 1, 1 }; + } + const bool reserved = demand != nullptr ? + common_reserve_activation_publication(context, dispatch_match, *publication_subject, *demand, + publication_name, publication) : + common_reserve_private_activation_output(context, dispatch_match, private_layout, + CommonActivationPublicationFormat::Q8_1X4, + "common.decode.flash_attention.private_q8_1_x4", publication); + if (!reserved || + (demand != nullptr && !common_append_activation_publication(dispatch_match, publication, publication_name))) { return false; } + const size_t q8_row_bytes = publication.byte_count / static_cast(match.query_token_count); const size_t query_row_bytes = static_cast(query_hidden_size) * sizeof(float); const size_t mask_row_bytes = static_cast(match.key_value_token_count) * sizeof(ggml_fp16_t); @@ -636,6 +643,7 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa for (int64_t row = 0; row < match.query_token_count; ++row) { Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kFlashAttentionDecodeSplitNextQ8Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; dispatch.kernel.integer_parameters.emplace("key_value_token_count", match.key_value_token_count); add_flash_attention_decode_compile_parameters(dispatch.kernel, match.query_head_count, match.key_value_head_count, match.qk_head_size, @@ -654,7 +662,7 @@ static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMa { completion_counter, 0, static_cast(match.key_value_head_count) * sizeof(int32_t) }); dispatch.bindings.push_back( { match.output->id, static_cast(row) * match.output->nb[2], output_row_bytes }); - dispatch.bindings.push_back({ q8_output, static_cast(row) * q8_row_bytes, q8_row_bytes }); + dispatch.bindings.push_back(publication.binding(static_cast(row) * q8_row_bytes, q8_row_bytes)); dispatch_match.dispatches.push_back(std::move(dispatch)); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp index 496938e8fb19..c0016ba9f7d4 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -1,5 +1,6 @@ #include "dispatch-gated-mul-mat.h" +#include "dispatch-activation-publication.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -19,14 +20,20 @@ static constexpr KernelCatalogRef kMulMatSwiGLUF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledPairAlignedF32BinaryPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned"); static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32K16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"); static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledPairAlignedF32BinaryBiasAddPublishF32Kernel = GGML_HRX_KERNEL_REF( + "loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32_aligned"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_decode_wave64"); static constexpr KernelCatalogRef kMulMatSwiGLUF32F32LowTokenDotKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_lowtoken_dot"); +static constexpr KernelCatalogRef kMulMatSwiGLUQ1_0Q8_1X4Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_q1_0_q8_1_x4_f32"); static constexpr KernelCatalogRef kMulMatSwiGLUQ4Q8LowTokenDotKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot"); static constexpr KernelCatalogRef kMulMatSwiGLUQ4Q8OutputKernel = @@ -350,11 +357,18 @@ static bool match_packed_mul_mat_glu_dispatch(const DispatchMatchContext & conte return false; } - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatTiledPairF32BinaryPublishF32Kernel); + const bool use_aligned_dynamic = match.token_count % 32 == 0; + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(use_aligned_dynamic ? kMulMatTiledPairAlignedF32BinaryPublishF32Kernel : + kMulMatTiledPairF32BinaryPublishF32Kernel); + if (use_aligned_dynamic) { + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + } dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", - common_to_config_value(match.token_count)); + if (!use_aligned_dynamic) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + } dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", @@ -624,13 +638,19 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte return false; } - const GraphNode * consumer = common_find_only_consumer_with_op(context.graph, match.output->id, GGML_OP_MUL_MAT); - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(context.graph, consumer, kMulMatSwiGLUQ4Q8PrefillKernel, false); - const bool packed_output = projection.matched() && projection.input->id == match.output->id && - projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, - projection.output_size, projection.token_count, true); + const bool has_only_projection = + common_find_only_consumer_with_op(context.graph, match.output->id, GGML_OP_MUL_MAT) != nullptr; + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputF16Row | + (has_only_projection ? DispatchActivationInputF16K16Major : DispatchActivationInputNone), + DispatchActivationInputNone, + DispatchActivationProducerRequirementPackedK16Economy, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + const CommonActivationPublicationDemand output_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + const bool packed_output = output_demand.format == CommonActivationPublicationFormat::F16K16Major; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUQ4Q8PrefillKernel); @@ -650,21 +670,21 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte match.input_size, match.output_size, weight->byte_count }); } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - const size_t f16_bytes = match.output->byte_count / 2; - const ValueId f16_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); const char * f16_name = packed_output ? "common.mul_mat_swiglu.k16_major_f16" : "common.mul_mat_swiglu.f16"; - dispatch_match.transients.push_back({ f16_output, f16_name, f16_bytes, 256 }); - dispatch.bindings.push_back({ f16_output, 0, f16_bytes }); - Status status; - const bool recorded = - packed_output ? - dispatch_match.metadata.append_generated_resource( - { match.output->id, GeneratedResourceRole::F16K16Major, f16_output, f16_bytes, {} }, status) : - dispatch_match.metadata.append_alternate_value( - { match.output->id, f16_output, GGML_TYPE_F16, f16_bytes, f16_name }, status); - if (!recorded) { - dispatch_match.status.append(status); - return false; + if (output_demand.matched()) { + CommonActivationPublication publication; + if (!common_reserve_activation_publication( + context, dispatch_match, *match.output, output_demand, f16_name, publication) || + !common_append_activation_publication(dispatch_match, publication, f16_name)) { + return false; + } + dispatch.bindings.push_back(publication.binding()); + } else { + const size_t f16_bytes = match.output->byte_count / 2; + const ValueId f16_output(context.next_plan_value.value + + static_cast(dispatch_match.transients.size())); + dispatch_match.transients.push_back({ f16_output, f16_name, f16_bytes, 256 }); + dispatch.bindings.push_back({ f16_output, 0, f16_bytes }); } for (const GraphNode * node : { match.gate_node, match.up_node, match.glu_node }) { if (!append_covered_node_index_once(context.graph, context.covered_nodes, node, dispatch_match.covered_nodes)) { @@ -681,30 +701,21 @@ static bool supports_swiglu_q8_output_shape(const MulMatSwiGLUMatch & match) { common_is_supported_dense_output_size(match.output_size) && match.output_size % 128 == 0; } -static bool has_qualified_swiglu_q8_consumer(const Graph & graph, const MulMatSwiGLUMatch & match) { - if (!supports_swiglu_q8_output_shape(match) || !distinct_storage(graph, *match.input, *match.gate_weight) || - !distinct_storage(graph, *match.input, *match.up_weight) || - !distinct_storage(graph, *match.input, *match.output) || - !distinct_storage(graph, *match.gate_weight, *match.up_weight) || - !distinct_storage(graph, *match.gate_weight, *match.output) || - !distinct_storage(graph, *match.up_weight, *match.output)) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { - CommonMulMatMatch projection = - common_match_mul_mat_any_format(graph, consumer, kMulMatSwiGLUQ4Q8OutputKernel, true); - if (!projection.matched()) { - projection = common_match_mul_mat_any_format(graph, consumer, kMulMatSwiGLUQ4Q8OutputKernel, false); - } - if (projection.matched() && projection.input->id == match.output->id && - projection.weight->alias_source.value < 0 && projection.output_size >= 64 && - projection.output_size % 64 == 0 && - (projection.weight_format == CommonMulMatWeightFormat::Q4K || - projection.weight_format == CommonMulMatWeightFormat::Q6K)) { - return true; - } - } - return false; +static CommonActivationPublicationDemand select_swiglu_q8_publication(const DispatchMatchContext & context, + const MulMatSwiGLUMatch & match) { + const bool eligible = supports_swiglu_q8_output_shape(match) && + distinct_storage(context.graph, *match.input, *match.gate_weight) && + distinct_storage(context.graph, *match.input, *match.up_weight) && + distinct_storage(context.graph, *match.input, *match.output) && + distinct_storage(context.graph, *match.gate_weight, *match.up_weight) && + distinct_storage(context.graph, *match.gate_weight, *match.output) && + distinct_storage(context.graph, *match.up_weight, *match.output); + const CommonActivationPublicationCapabilities capabilities = { + eligible ? DispatchActivationInputQ8_1X4 : DispatchActivationInputNone, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + return publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; } static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLUMatch & match) { @@ -721,22 +732,20 @@ static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLU return !use_q8 && !use_direct_dot && match.token_count >= 2; } -static bool has_qualified_swiglu_k16_consumer(const Graph & graph, const MulMatSwiGLUMatch & match) { - if (!graph.has_index() || match.op != BinaryKind::SwiGLU || match.output == nullptr || - match.output->ne[1] != match.token_count || match.output->ne[2] != 1 || match.output->ne[3] != 1) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(graph, consumer, kMulMatTiledPairF32BinaryPublishF32Kernel, false); - if (projection.matched() && projection.input->id == match.output->id && - projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, projection.output_size, - projection.token_count, true)) { - return true; - } +static CommonActivationPublicationDemand select_swiglu_k16_publication(const DispatchMatchContext & context, + const MulMatSwiGLUMatch & match) { + if (match.op != BinaryKind::SwiGLU || match.output == nullptr || match.output->ne[1] != match.token_count || + match.output->ne[2] != 1 || match.output->ne[3] != 1) { + return {}; } - return false; + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputF16K16Major, + DispatchActivationInputNone, + DispatchActivationProducerRequirementPackedK16Economy, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + return publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; } static MulMatSwiGLUPostOpsMatch match_mul_mat_swiglu_postops(const DispatchMatchContext & context) { @@ -799,22 +808,35 @@ static MulMatSwiGLUPostOpsMatch match_mul_mat_swiglu_postops(const DispatchMatch return match; } -static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, +static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, const MulMatSwiGLUPostOpsMatch * postops, - DispatchMatch & dispatch_match) { + DispatchMatch & dispatch_match) { const MulMatSwiGLUMatch & match = postops != nullptr ? postops->root : match_mul_mat_swiglu(context); if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { return false; } - const bool publish_k16 = postops == nullptr && has_qualified_swiglu_k16_consumer(context.graph, match); - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(postops != nullptr ? postops->kernel : - publish_k16 ? kMulMatTiledPairF32BinaryPublishF32K16Kernel : - kMulMatTiledPairF32BinaryPublishF32Kernel); + const CommonActivationPublicationDemand k16_demand = + postops == nullptr ? select_swiglu_k16_publication(context, match) : CommonActivationPublicationDemand{}; + const bool publish_k16 = k16_demand.matched(); + const bool use_aligned_dynamic = match.token_count % 32 == 0 && !publish_k16; + const KernelCatalogRef selected_kernel = + use_aligned_dynamic ? + (postops != nullptr ? kMulMatTiledPairAlignedF32BinaryBiasAddPublishF32Kernel : + kMulMatTiledPairAlignedF32BinaryPublishF32Kernel) : + (postops != nullptr ? postops->kernel : + publish_k16 ? kMulMatTiledPairF32BinaryPublishF32K16Kernel : + kMulMatTiledPairF32BinaryPublishF32Kernel); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(selected_kernel); + if (use_aligned_dynamic) { + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + } dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", - common_to_config_value(match.token_count)); + if (!use_aligned_dynamic) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + } dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", @@ -850,16 +872,13 @@ static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext dispatch.bindings.push_back({ output->id, 0, output->byte_count }); if (publish_k16) { - const size_t bytes = output->byte_count / 2; - const ValueId packed(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); - Status status; - if (!dispatch_match.metadata.append_generated_resource( - { output->id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, status)) { - dispatch_match.status.append(status); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, dispatch_match, *output, k16_demand, + "common.mul_mat_swiglu.k16_major_f16", publication) || + !common_append_activation_publication(dispatch_match, publication, "common.mul_mat_swiglu.k16_major_f16")) { return false; } - dispatch.bindings.push_back({ packed, 0, bytes }); - dispatch_match.transients.push_back({ packed, "common.mul_mat_swiglu.k16_major_f16", bytes, 256 }); + dispatch.bindings.push_back(publication.binding()); } if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, @@ -905,12 +924,48 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, return false; } + const bool use_q1_0 = match.token_count == 1 && match.input_size % 128 == 0 && + match.op == BinaryKind::SwiGLU && match.gate_weight->alias_source.value < 0 && + match.up_weight->alias_source.value < 0 && + match.gate_format == CommonMulMatWeightFormat::Q1_0 && + match.up_format == CommonMulMatWeightFormat::Q1_0 && + distinct_storage(context.graph, *match.gate_weight, *match.up_weight); + if (use_q1_0) { + DispatchBinding activation; + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUQ1_0Q8_1X4Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.q1_0_q8_1_x4.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.q1_0_q8_1_x4.output_size", + common_to_config_value(match.output_size)); + dispatch.bindings.push_back(activation); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.up_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.glu_node, + dispatch_match.covered_nodes)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; + } + } + const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 && match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 && match.gate_format == CommonMulMatWeightFormat::Q4K && match.up_format == CommonMulMatWeightFormat::Q4K; const bool use_q8 = pack_q4; - const bool publish_q8 = use_q8 && has_qualified_swiglu_q8_consumer(context.graph, match); + const CommonActivationPublicationDemand q8_demand = + use_q8 ? select_swiglu_q8_publication(context, match) : CommonActivationPublicationDemand{}; + const bool publish_q8 = q8_demand.matched(); DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; if (use_q8 && !common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, @@ -956,17 +1011,14 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_q8) { - const ValueId q8_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); - const size_t bytes = static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.output_size); constexpr const char * name = "common.mul_mat_swiglu.q8_1_x4"; - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { match.output->id, q8_output, GGML_TYPE_Q8_1, bytes, name }, status)) { - dispatch_match.status.append(status); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, q8_demand, name, + publication) || + !common_append_activation_publication(dispatch_match, publication, name)) { return false; } - dispatch.bindings.push_back({ q8_output, 0, bytes }); - dispatch_match.transients.push_back({ q8_output, name, bytes, 256 }); + dispatch.bindings.push_back(publication.binding()); } if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, 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 73b4ca960656..022d1f05279e 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 @@ -1,11 +1,12 @@ #include "dispatch-get-rows.h" #include "../qwen/dispatch-llm-profiles.h" +#include "dispatch-activation-publication.h" +#include "dispatch-mul-mat-common.h" #include "dispatch-mul-mat-weight-format.h" #include "ggml.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" -#include #include #include #include @@ -25,8 +26,6 @@ static constexpr KernelCatalogRef kGetRowsRmsNormBinaryQ8_1X4F16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); static constexpr int64_t kMaximumHiddenElements = int64_t{ 1 } << 30; static constexpr int64_t kQ1_0GetRowsFormat = 10; -static constexpr int64_t kQwenHiddenSize = kQwen30BMoeDispatchProfile.hidden_size; -static constexpr int64_t kQwenVocabularyCount = 151936; static constexpr int64_t kMaxGetRowsRowCount = 262208; static const Value * graph_value(const Graph & graph, ValueId id) { @@ -54,6 +53,14 @@ static bool is_supported_token_count(int64_t token_count) { return token_count >= 1 && token_count <= 2048; } +static int64_t token_capacity_class(int64_t token_count) { + int64_t capacity = 1; + while (capacity < token_count) { + capacity *= 2; + } + return capacity; +} + static bool is_supported_row_count(int64_t row_count) { return row_count >= 1 && row_count <= kMaxGetRowsRowCount; } @@ -69,66 +76,6 @@ static std::string to_config_value(float value) { return out.str(); } -static size_t row_byte_count(ggml_type type, int64_t token_count, int64_t hidden_size) { - if (token_count <= 0 || hidden_size <= 0) { - return 0; - } - return static_cast(token_count) * ggml_row_size(type, hidden_size); -} - -static bool is_qwen_q6k_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { - if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || - consumer->inputs[1] != input.id) { - return false; - } - - const Value * weight = graph_value(graph, consumer->inputs[0]); - const Value * output = graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || weight->type != GGML_TYPE_Q6_K || output->type != GGML_TYPE_F32 || - !weight->contiguous || !output->contiguous) { - return false; - } - - return input.ne[0] == kQwenHiddenSize && input.ne[1] == 1 && input.ne[2] == 1 && input.ne[3] == 1 && - weight->ne[0] == kQwenHiddenSize && weight->ne[1] == kQwenVocabularyCount && weight->ne[2] == 1 && - weight->ne[3] == 1 && output->ne[0] == kQwenVocabularyCount && output->ne[1] == 1 && output->ne[2] == 1 && - output->ne[3] == 1; -} - -static void append_unique_demand(std::vector & demands, ggml_type type) { - if (std::find(demands.begin(), demands.end(), type) == demands.end()) { - demands.push_back(type); - } -} - -static std::vector collect_alternate_demands(const Graph & graph, const Value & value) { - std::vector demands; - if (!graph.has_index()) { - return demands; - } - - const std::vector & consumers = graph.index().consumers(value.id); - for (const GraphNode * consumer : consumers) { - if (is_qwen_q6k_q8_consumer(graph, consumer, value)) { - append_unique_demand(demands, GGML_TYPE_Q8_1); - } - } - return demands; -} - -static const char * alternate_name(ggml_type type) { - switch (type) { - case GGML_TYPE_Q8_1: - return "common.get_rows.q8_1_x4"; - case GGML_TYPE_F16: - return "common.get_rows.f16"; - case GGML_TYPE_F32: - return "common.get_rows.f32"; - default: - return "common.get_rows.next"; - } -} - static bool get_rows_format_for_type(ggml_type type, int64_t & format) { if (type == GGML_TYPE_Q1_0) { format = kQ1_0GetRowsFormat; @@ -356,8 +303,13 @@ static GetRowsRmsNormMatch match_get_rows_rmsnorm_binary(const DispatchMatchCont static Dispatch make_get_rows_dispatch(const GetRowsMatch & match) { Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kGetRowsF32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; add_common_integer_parameters(dispatch, match); - add_common_compile_parameters(dispatch, match); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.token_capacity", + to_config_value(token_capacity_class(match.token_count))); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.hidden_capacity", to_config_value(match.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.weight_format", + to_config_value(match.weight_format_value)); add_primary_bindings(dispatch, match); return dispatch; } @@ -377,20 +329,6 @@ static Dispatch make_get_rows_next_dispatch(const GetRowsMatch & match, return dispatch; } -static bool append_alternate_metadata(DispatchMatch & dispatch_match, - const Value & output, - ValueId alternate_value, - ggml_type type, - size_t byte_count) { - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { output.id, alternate_value, type, byte_count, alternate_name(type) }, metadata_status)) { - dispatch_match.status.append(metadata_status); - return false; - } - 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()) { @@ -433,42 +371,33 @@ static bool match_get_rows_f32_next_dispatch(const DispatchMatchContext & contex return false; } - const std::vector demands = collect_alternate_demands(context.graph, *match.output); - if (demands.empty()) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputNone, + DispatchActivationProducerRequirementNone, + DispatchActivationInputNone, + DispatchActivationConsumerSourceQwen, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + if (!publication_plan.matched()) { return false; } + const CommonActivationPublicationDemand & demand = publication_plan.preferred(); - for (ggml_type type : demands) { - CommonMulMatWeightFormat next_format; - if (!common_mul_mat_alternate_format_for_type(type, next_format)) { - return false; - } - - const size_t next_byte_count = row_byte_count(type, match.token_count, match.hidden_size); - if (next_byte_count == 0) { - return false; - } - - if (type == GGML_TYPE_F32 && next_byte_count == match.output->byte_count) { - if (!append_alternate_metadata(dispatch_match, *match.output, match.output->id, type, next_byte_count)) { - return false; - } - continue; - } - - const ValueId next_value(context.next_plan_value.value + - static_cast(dispatch_match.transients.size())); - dispatch_match.dispatches.push_back( - make_get_rows_next_dispatch(match, next_format, next_value, next_byte_count)); - dispatch_match.transients.push_back({ next_value, alternate_name(type), next_byte_count, 256 }); - if (!append_alternate_metadata(dispatch_match, *match.output, next_value, type, next_byte_count)) { - return false; - } + CommonMulMatWeightFormat next_format; + if (!common_mul_mat_alternate_format_for_type(GGML_TYPE_Q8_1, next_format)) { + return false; } - - if (dispatch_match.dispatches.empty()) { - dispatch_match.dispatches.push_back(make_get_rows_dispatch(match)); + CommonActivationPublication publication; + constexpr const char * publication_name = "common.get_rows.q8_1_x4"; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, demand, + publication_name, publication) || + !common_append_activation_publication(dispatch_match, publication, publication_name)) { + return false; } + dispatch_match.dispatches.push_back( + make_get_rows_next_dispatch(match, next_format, publication.alternate_value, publication.byte_count)); dispatch_match.covered_nodes.push_back(context.root_index); return dispatch_match.status.success(); } @@ -482,13 +411,43 @@ static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & return false; } - const size_t q8_bytes = row_byte_count(GGML_TYPE_Q8_1, match.rows.token_count, match.rows.hidden_size); - const size_t f16_bytes = row_byte_count(GGML_TYPE_F16, match.rows.token_count, match.rows.hidden_size); - if (q8_bytes == 0 || f16_bytes == 0) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4 | DispatchActivationInputF16Row, + DispatchActivationInputQ8_1X4 | DispatchActivationInputF16Row, + DispatchActivationProducerRequirementNone, + DispatchActivationInputQ8_1X4 | DispatchActivationInputF16Row, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.normalized_output, capabilities); + const CommonActivationPublicationDemand * q8_demand = + publication_plan.find(CommonActivationPublicationFormat::Q8_1X4); + const CommonActivationPublicationDemand * f16_demand = + publication_plan.find(CommonActivationPublicationFormat::F16Row); + const bool publish = publication_plan.matched(); + if (publish && (q8_demand == nullptr || f16_demand == nullptr)) { + return false; + } + CommonActivationPublication q8_publication; + CommonActivationPublication f16_publication; + const bool q8_reserved = publish ? + common_reserve_activation_publication(context, dispatch_match, *match.normalized_output, *q8_demand, + "common.get_rows_rmsnorm.q8_1_x4", q8_publication) : + common_reserve_private_activation_output(context, dispatch_match, *match.normalized_output, + CommonActivationPublicationFormat::Q8_1X4, + "common.get_rows_rmsnorm.private_q8_1_x4", q8_publication); + const bool f16_reserved = publish ? + common_reserve_activation_publication(context, dispatch_match, *match.normalized_output, *f16_demand, + "common.get_rows_rmsnorm.f16", f16_publication) : + common_reserve_private_activation_output(context, dispatch_match, *match.normalized_output, + CommonActivationPublicationFormat::F16Row, + "common.get_rows_rmsnorm.private_f16", f16_publication); + if (!q8_reserved || !f16_reserved || + (publish && (!common_append_activation_publication(dispatch_match, q8_publication, + "common.get_rows_rmsnorm.q8_1_x4") || + !common_append_activation_publication(dispatch_match, f16_publication, + "common.get_rows_rmsnorm.f16")))) { return false; } - const ValueId q8_value = context.next_plan_value; - const ValueId f16_value(context.next_plan_value.value + 1); Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kGetRowsRmsNormBinaryQ8_1X4F16Kernel); @@ -500,25 +459,12 @@ static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & dispatch.bindings.push_back({ match.norm_weight->id, 0, match.norm_weight->byte_count }); dispatch.bindings.push_back({ match.rows.output->id, 0, match.rows.output->byte_count }); dispatch.bindings.push_back({ match.normalized_output->id, 0, match.normalized_output->byte_count }); - dispatch.bindings.push_back({ q8_value, 0, q8_bytes }); - dispatch.bindings.push_back({ f16_value, 0, f16_bytes }); - - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { match.normalized_output->id, q8_value, GGML_TYPE_Q8_1, q8_bytes, "common.get_rows_rmsnorm.q8_1_x4" }, - status) || - !dispatch_match.metadata.append_alternate_value( - { match.normalized_output->id, f16_value, GGML_TYPE_F16, f16_bytes, "common.get_rows_rmsnorm.f16" }, - status)) { - dispatch_match.status.append(status); - return false; - } + dispatch.bindings.push_back(q8_publication.binding()); + dispatch.bindings.push_back(f16_publication.binding()); dispatch_match.covered_nodes.push_back(context.root_index); dispatch_match.covered_nodes.push_back(match.rms_node_index); dispatch_match.covered_nodes.push_back(match.binary_node_index); - dispatch_match.transients.push_back({ q8_value, "common.get_rows_rmsnorm.q8_1_x4", q8_bytes, 256 }); - dispatch_match.transients.push_back({ f16_value, "common.get_rows_rmsnorm.f16", f16_bytes, 256 }); dispatch_match.dispatches.push_back(std::move(dispatch)); return dispatch_match.status.success(); } 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 25a024d7b01c..93127a1e6c58 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 @@ -3,6 +3,7 @@ #include "../dispatch-registry.h" #include "dispatch-binary-common.h" #include "dispatch-mul-mat-weight-format.h" +#include "dispatch-symmetric-i4.h" #include "ggml.h" #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" @@ -10,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -184,6 +186,16 @@ inline ggml_type common_mul_mat_format_type(CommonMulMatWeightFormat format) { return GGML_TYPE_Q5_0; case CommonMulMatWeightFormat::Q5_1: return GGML_TYPE_Q5_1; + case CommonMulMatWeightFormat::IQ1_S: + return GGML_TYPE_IQ1_S; + case CommonMulMatWeightFormat::IQ1_M: + return GGML_TYPE_IQ1_M; + case CommonMulMatWeightFormat::IQ2_XXS: + return GGML_TYPE_IQ2_XXS; + case CommonMulMatWeightFormat::IQ2_XS: + return GGML_TYPE_IQ2_XS; + case CommonMulMatWeightFormat::IQ3_XXS: + return GGML_TYPE_IQ3_XXS; case CommonMulMatWeightFormat::IQ2_S: return GGML_TYPE_IQ2_S; case CommonMulMatWeightFormat::IQ3_S: @@ -302,41 +314,14 @@ inline const GraphNode * common_find_only_consumer_with_op(const Graph & graph, return consumers.front(); } -inline bool common_has_direct_symmetric_i4_lowrow_consumer(const Graph & graph, const Value & value) { - if (!graph.has_index() || value.type != GGML_TYPE_F32 || !value.contiguous || value.ne[0] < 256 || - value.ne[0] > 32768 || value.ne[0] % 64 != 0 || value.element_count <= 0 || - value.element_count % value.ne[0] != 0) { - return false; - } - - const int64_t token_count = value.element_count / value.ne[0]; - if (token_count < 1 || token_count > 16) { +inline bool common_has_symmetric_i4_lowrow_consumer(const Graph & graph, const Value & value) { + if (!graph.has_index()) { return false; } - for (const GraphNode * consumer : graph.index().consumers(value.id)) { - if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || - consumer->inputs[1] != value.id) { - continue; - } - - const Value * weight = common_graph_value(graph, consumer->inputs[0]); - const Value * output = common_graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || - (weight->type != GGML_TYPE_Q5_K && weight->type != GGML_TYPE_IQ4_XS) || weight->alias_source.value >= 0 || - !weight->contiguous || !output->contiguous || output->type != GGML_TYPE_F32 || - weight->ne[0] != value.ne[0] || weight->ne[1] != output->ne[0] || output->ne[0] % 64 != 0 || - output->ne[1] != token_count || output->ne[2] != 1 || output->ne[3] != 1) { - continue; + if (consumer != nullptr && common_symmetric_i4_lowrow_mul_mat_eligible(graph, *consumer, value)) { + return true; } - return true; - } - return false; -} - -inline bool common_has_symmetric_i4_lowrow_consumer(const Graph & graph, const Value & value) { - if (common_has_direct_symmetric_i4_lowrow_consumer(graph, value)) { - return true; } for (const GraphNode * consumer : graph.index().consumers(value.id)) { @@ -348,9 +333,13 @@ inline bool common_has_symmetric_i4_lowrow_consumer(const Graph & graph, const V const Value * reshaped = common_graph_value(graph, consumer->output); if (reshaped != nullptr && reshaped->type == GGML_TYPE_F32 && reshaped->contiguous && reshaped->storage_root == value.storage_root && reshaped->element_count == value.element_count && - reshaped->byte_count == value.byte_count && - common_has_direct_symmetric_i4_lowrow_consumer(graph, *reshaped)) { - return true; + reshaped->byte_count == value.byte_count) { + for (const GraphNode * reshaped_consumer : graph.index().consumers(reshaped->id)) { + if (reshaped_consumer != nullptr && + common_symmetric_i4_lowrow_mul_mat_eligible(graph, *reshaped_consumer, *reshaped)) { + return true; + } + } } } return false; @@ -402,6 +391,20 @@ inline CommonMulMatMatch common_match_mul_mat_any_format(const Graph & graph CommonMulMatWeightFormat format = CommonMulMatWeightFormat::Q4K; if (!common_mul_mat_format_for_type(weight->type, format)) { + if (weight->type == GGML_TYPE_IQ1_S) { + format = CommonMulMatWeightFormat::IQ1_S; + } else if (weight->type == GGML_TYPE_IQ1_M) { + format = CommonMulMatWeightFormat::IQ1_M; + } else { + return {}; + } + } + if ((format == CommonMulMatWeightFormat::IQ1_S || format == CommonMulMatWeightFormat::IQ1_M) && + (std::getenv("GGML_HRX_DISABLE_IQ_CODEBOOK_MATMUL") != nullptr || + (format == CommonMulMatWeightFormat::IQ1_S && + std::getenv("GGML_HRX_DISABLE_IQ1_S_CODEBOOK_MATMUL") != nullptr) || + (format == CommonMulMatWeightFormat::IQ1_M && + std::getenv("GGML_HRX_DISABLE_IQ1_M_CODEBOOK_MATMUL") != nullptr))) { return {}; } @@ -428,6 +431,118 @@ inline CommonMulMatMatch common_match_mul_mat_any_format(const Graph & graph return match; } +struct CommonMulMatActivationConsumerMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + CommonMulMatWeightFormat weight_format = CommonMulMatWeightFormat::Q4K; + + bool matched() const { return input != nullptr && weight != nullptr && output != nullptr; } +}; + +inline CommonMulMatActivationConsumerMatch common_match_mul_mat_activation_consumer(const Graph & graph, + const GraphNode * consumer, + const Value & input) { + if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || + consumer->inputs[1] != input.id) { + return {}; + } + + const Value * weight = common_graph_value(graph, consumer->inputs[0]); + const Value * output = common_graph_value(graph, consumer->output); + if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input.contiguous || !weight->contiguous || !output->contiguous || !common_is_2d(input) || + !common_is_2d(*weight) || !common_is_2d(*output)) { + return {}; + } + + CommonMulMatWeightFormat format; + if (!common_mul_mat_format_for_type(weight->type, format)) { + return {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input.ne[1]; + if (input.ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count || + !common_is_supported_dense_input_size(input_size) || !common_is_supported_dense_output_size(output_size) || + output_size % 64 != 0) { + return {}; + } + + return { &input, weight, output, input_size, output_size, token_count, format }; +} + +inline bool common_accepts_q8_1_x4_decode_mul_mat(const Graph & graph, + const GraphNode * consumer, + const Value & input) { + const CommonMulMatActivationConsumerMatch match = + common_match_mul_mat_activation_consumer(graph, consumer, input); + return match.matched() && match.weight->alias_source.value < 0 && match.token_count >= 1 && + match.token_count <= 5 && + (match.weight_format == CommonMulMatWeightFormat::Q4K || + match.weight_format == CommonMulMatWeightFormat::Q6K); +} + +inline bool common_accepts_q8_1_x4_decode_mul_mat(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return common_accepts_q8_1_x4_decode_mul_mat(context.graph, &consumer, input); +} + +inline bool common_accepts_q8_1_x4_prefill_mul_mat(const Graph & graph, + const GraphNode * consumer, + const Value & input) { + const CommonMulMatActivationConsumerMatch match = + common_match_mul_mat_activation_consumer(graph, consumer, input); + return match.matched() && match.token_count >= 256 && match.token_count <= 2048 && + match.token_count % 256 == 0 && + (match.weight_format == CommonMulMatWeightFormat::Q5K || + match.weight_format == CommonMulMatWeightFormat::IQ4_XS); +} + +inline bool common_accepts_f16_k16_major_mul_mat(const Graph & graph, + const GraphNode * consumer, + const Value & input, + bool packed_producer = false) { + const CommonMulMatActivationConsumerMatch match = + common_match_mul_mat_activation_consumer(graph, consumer, input); + return match.matched() && match.weight->alias_source.value < 0 && + common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, + match.token_count, packed_producer); +} + +inline bool common_accepts_f16_k16_major_mul_mat(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return common_accepts_f16_k16_major_mul_mat(context.graph, &consumer, input); +} + +inline bool common_accepts_f16_k16_major_packed_producer_mul_mat(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return common_accepts_f16_k16_major_mul_mat(context.graph, &consumer, input, true); +} + +inline bool common_accepts_f16_row_prefill_mul_mat(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + const CommonMulMatActivationConsumerMatch match = + common_match_mul_mat_activation_consumer(context.graph, &consumer, input); + if (!match.matched() || match.weight->alias_source.value >= 0 || + !common_is_supported_prefill_token_count(match.token_count)) { + return false; + } + if (match.weight_format == CommonMulMatWeightFormat::Q6K && match.token_count % 128 == 0) { + return true; + } + return match.weight_format == CommonMulMatWeightFormat::Q4K && match.token_count % 256 == 0 && + match.output_size >= match.input_size / 4; +} + // Share the conversion used by F16-operand matmuls. inline bool common_prepare_f16_input(const DispatchMatchContext & context, const Value & input, 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 492450ae1202..cc3848c42e10 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 @@ -18,6 +18,11 @@ enum class CommonMulMatWeightFormat { Q4_1, Q5_0, Q5_1, + IQ1_S, + IQ1_M, + IQ2_XXS, + IQ2_XS, + IQ3_XXS, IQ2_S, IQ3_S, IQ4_NL, @@ -58,6 +63,15 @@ inline bool common_mul_mat_format_for_type(ggml_type type, CommonMulMatWeightFor case GGML_TYPE_Q5_1: format = CommonMulMatWeightFormat::Q5_1; return true; + case GGML_TYPE_IQ2_XXS: + format = CommonMulMatWeightFormat::IQ2_XXS; + return true; + case GGML_TYPE_IQ2_XS: + format = CommonMulMatWeightFormat::IQ2_XS; + return true; + case GGML_TYPE_IQ3_XXS: + format = CommonMulMatWeightFormat::IQ3_XXS; + return true; case GGML_TYPE_IQ2_S: format = CommonMulMatWeightFormat::IQ2_S; return true; @@ -114,6 +128,16 @@ inline int64_t common_mul_mat_format_config_value(CommonMulMatWeightFormat forma return 50; case CommonMulMatWeightFormat::Q5_1: return 51; + case CommonMulMatWeightFormat::IQ1_S: + return 19; + case CommonMulMatWeightFormat::IQ1_M: + return 29; + case CommonMulMatWeightFormat::IQ2_XXS: + return 24; + case CommonMulMatWeightFormat::IQ2_XS: + return 25; + case CommonMulMatWeightFormat::IQ3_XXS: + return 18; case CommonMulMatWeightFormat::IQ2_S: return 22; case CommonMulMatWeightFormat::IQ3_S: 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 5a35e811c092..047caa6acec3 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,8 +1,11 @@ #include "dispatch-mul-mat.h" +#include "dispatch-activation-publication.h" #include "dispatch-mul-mat-common.h" +#include "dispatch_registration/immutable-program-resource.h" #include "ggml.h" #include "graph/graph-matcher.h" +#include "iq-codebook-resource.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include @@ -18,6 +21,20 @@ static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); static constexpr KernelCatalogRef kMulMatTiledF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledIQ1SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledIQ1MF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledIQ2SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledIQ3SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledAlignedF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_aligned"); +static constexpr KernelCatalogRef kMulMatTiledToken64AlignedF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_token64_aligned"); +static constexpr KernelCatalogRef kMulMatTiledIQ3XXSF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32"); static constexpr KernelCatalogRef kMulMatTiledF32F16AlternateKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate"); static constexpr KernelCatalogRef kMulMatSkinnyF32F32Kernel = @@ -30,12 +47,28 @@ static constexpr KernelCatalogRef kMulMatAddF32F32DecodeWave64Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_decode_wave64"); static constexpr KernelCatalogRef kMulMatVectorF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorQ1_0Q8_1X4F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q1_0_q8_1_x4_f32"); +static constexpr KernelCatalogRef kMulMatVectorQ1_0Q8_1X4Rows2F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q1_0_q8_1_x4_rows2_f32"); +static constexpr KernelCatalogRef kMulMatVectorIQ1SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_iq1_s_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorIQ1MF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_iq1_m_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorIQ2SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_iq2_s_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorIQ3SF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_iq3_s_f32_f32"); +static constexpr KernelCatalogRef kMulMatVectorIQ3XXSF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_iq3_xxs_f32_f32"); static constexpr KernelCatalogRef kMulMatVectorBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_bias_residual_f32_f32"); static constexpr KernelCatalogRef kQuantizeQ8_1X4F32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); static constexpr KernelCatalogRef kMulMatTiledBiasAddF32F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatTiledAlignedBiasAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32_aligned"); static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K64PlaneKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i4_k64_plane"); static constexpr KernelCatalogRef kMulMatSymmetricI4WmmaKernel = @@ -87,6 +120,97 @@ static bool deprecated_mul_mat_dispatch_disabled() { return std::getenv("GGML_HRX_DISABLE_DEPRECATED_MUL_MAT_DISPATCH") != nullptr; } +static bool iq_codebook_matmul_disabled(CommonMulMatWeightFormat format) { + if (std::getenv("GGML_HRX_DISABLE_IQ_CODEBOOK_MATMUL") != nullptr) { + return true; + } + if (format == CommonMulMatWeightFormat::IQ2_XXS) { + return std::getenv("GGML_HRX_DISABLE_IQ2_XXS_CODEBOOK_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ2_XS) { + return std::getenv("GGML_HRX_DISABLE_IQ2_XS_CODEBOOK_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ2_S) { + return std::getenv("GGML_HRX_DISABLE_IQ2_S_GRID_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ3_S) { + return std::getenv("GGML_HRX_DISABLE_IQ3_S_CODEBOOK_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ3_XXS) { + return std::getenv("GGML_HRX_DISABLE_IQ3_XXS_CODEBOOK_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ1_S) { + return std::getenv("GGML_HRX_DISABLE_IQ1_S_CODEBOOK_MATMUL") != nullptr; + } + if (format == CommonMulMatWeightFormat::IQ1_M) { + return std::getenv("GGML_HRX_DISABLE_IQ1_M_CODEBOOK_MATMUL") != nullptr; + } + return false; +} + +static bool iq_codebook_matmul_required(CommonMulMatWeightFormat format) { + return format == CommonMulMatWeightFormat::IQ1_S || format == CommonMulMatWeightFormat::IQ1_M || + format == CommonMulMatWeightFormat::IQ2_XXS || format == CommonMulMatWeightFormat::IQ2_XS || + format == CommonMulMatWeightFormat::IQ3_XXS; +} + +static const IQCodebookResourceSpec * iq_codebook_for_weight_format(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::IQ1_S: + case CommonMulMatWeightFormat::IQ1_M: + return &iq_codebook_resource_spec(IQCodebookResource::IQ1); + case CommonMulMatWeightFormat::IQ2_XXS: + return &iq_codebook_resource_spec(IQCodebookResource::IQ2XXS); + case CommonMulMatWeightFormat::IQ2_XS: + return &iq_codebook_resource_spec(IQCodebookResource::IQ2XS); + case CommonMulMatWeightFormat::IQ3_XXS: + return &iq_codebook_resource_spec(IQCodebookResource::IQ3XXS); + case CommonMulMatWeightFormat::IQ2_S: + return &iq_codebook_resource_spec(IQCodebookResource::IQ2S); + case CommonMulMatWeightFormat::IQ3_S: + return &iq_codebook_resource_spec(IQCodebookResource::IQ3S); + default: + return nullptr; + } +} + +static ValueId require_iq_codebook(const DispatchMatchContext & context, + DispatchMatch & dispatch_match, + const IQCodebookResourceSpec & codebook) { + return require_immutable_program_resource(context, dispatch_match, + { codebook.key, codebook.data, codebook.alignment }); +} + +static KernelCatalogRef tiled_iq_codebook_kernel(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::IQ1_S: + return kMulMatTiledIQ1SF32F32Kernel; + case CommonMulMatWeightFormat::IQ1_M: + return kMulMatTiledIQ1MF32F32Kernel; + case CommonMulMatWeightFormat::IQ3_XXS: + return kMulMatTiledIQ3XXSF32F32Kernel; + case CommonMulMatWeightFormat::IQ3_S: + return kMulMatTiledIQ3SF32F32Kernel; + default: + return kMulMatTiledIQ2SF32F32Kernel; + } +} + +static KernelCatalogRef vector_iq_codebook_kernel(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::IQ1_S: + return kMulMatVectorIQ1SF32F32Kernel; + case CommonMulMatWeightFormat::IQ1_M: + return kMulMatVectorIQ1MF32F32Kernel; + case CommonMulMatWeightFormat::IQ3_XXS: + return kMulMatVectorIQ3XXSF32F32Kernel; + case CommonMulMatWeightFormat::IQ3_S: + return kMulMatVectorIQ3SF32F32Kernel; + default: + return kMulMatVectorIQ2SF32F32Kernel; + } +} + struct MulMatPostOpsMatch { const Value * input = nullptr; const Value * weight = nullptr; @@ -124,18 +248,9 @@ struct DecodeMulMatAddMatch { } }; -enum class VectorPublishFormat { - None, - F16, -}; - struct VectorPublishPlan { - VectorPublishFormat fused_format = VectorPublishFormat::None; - ValueId fused_output = {}; - size_t fused_byte_count = 0; - bool publish_q8 = false; - ValueId q8_output = {}; - size_t q8_byte_count = 0; + CommonActivationPublicationDemand demand; + CommonActivationPublication publication; }; static bool common_mul_mat_is_low_token_batch(int64_t token_count) { @@ -185,52 +300,60 @@ static bool is_bias_shape(const Value & value, int64_t output_size) { return true; } -static bool common_mul_mat_has_f16_alternate_consumer(const Graph & graph, const CommonMulMatMatch & match) { - if (!graph.has_index() || match.output == nullptr || match.output->type != GGML_TYPE_F32 || - !match.output->contiguous || match.token_count < 128) { +static bool accepts_tiled_q8_publication(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return common_accepts_q8_1_x4_prefill_mul_mat(context.graph, &consumer, input); +} + +static bool accepts_preferred_q8_publication(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + if (common_accepts_q8_1_x4_prefill_mul_mat(context.graph, &consumer, input)) { + return true; + } + if (!common_accepts_q8_1_x4_decode_mul_mat(context.graph, &consumer, input)) { return false; } + const Value * weight = context.graph.values().find(consumer.inputs[0]); + const Value * output = context.graph.values().find(consumer.output); + return weight != nullptr && output != nullptr && + (weight->type == GGML_TYPE_Q4_K || !context.graph.index().consumers(output->id).empty()); +} - for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { - const CommonMulMatMatch downstream = - common_match_mul_mat_any_format(graph, consumer, kMulMatQ4KF16WmmaPrefillWave32Kernel, false); - if (!downstream.matched() || downstream.input->id != match.output->id || - downstream.input_size != match.output_size || downstream.token_count != match.token_count || - downstream.output_size % 64 != 0 || downstream.input_size % 256 != 0) { - continue; - } - if (downstream.weight->type == GGML_TYPE_Q6_K && downstream.weight->alias_source.value < 0 && - downstream.token_count % 128 == 0) { - return true; - } - if (downstream.weight->type == GGML_TYPE_Q4_K && downstream.weight->alias_source.value < 0 && - downstream.token_count % 256 == 0 && downstream.output_size >= downstream.input_size / 4) { - return true; - } - } - return false; +static bool is_q1_0_q8_1_x4_vector_consumer(const CommonMulMatActivationConsumerMatch & match) { + return match.matched() && match.weight->alias_source.value < 0 && match.token_count == 1 && + match.input_size % 128 == 0 && match.weight_format == CommonMulMatWeightFormat::Q1_0; } -static bool common_mul_mat_has_q8_alternate_consumer(const Graph & graph, const CommonMulMatMatch & match) { - if (!graph.has_index() || match.output == nullptr || match.output->type != GGML_TYPE_F32 || - !match.output->contiguous || match.token_count < 256 || match.token_count > 2048 || - match.token_count % 256 != 0 || match.output_size % 128 != 0) { +static bool accepts_preferred_q1_0_q8_publication(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + const CommonMulMatActivationConsumerMatch match = + common_match_mul_mat_activation_consumer(context.graph, &consumer, input); + if (!is_q1_0_q8_1_x4_vector_consumer(match)) { return false; } - - for (const GraphNode * consumer : graph.index().consumers(match.output->id)) { - const CommonMulMatMatch downstream = - common_match_mul_mat_any_format(graph, consumer, kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel, false); - if (!downstream.matched() || downstream.input->id != match.output->id || - downstream.input_size != match.output_size || downstream.token_count != match.token_count || - downstream.output_size % 64 != 0 || downstream.input_size % 256 != 0) { + int64_t compatible_consumer_count = 0; + int64_t total_output_size = 0; + for (const GraphNode * candidate : context.graph.index().consumers(input.id)) { + const CommonMulMatActivationConsumerMatch candidate_match = + common_match_mul_mat_activation_consumer(context.graph, candidate, input); + if (!is_q1_0_q8_1_x4_vector_consumer(candidate_match)) { continue; } - if (downstream.weight->type == GGML_TYPE_Q5_K || downstream.weight->type == GGML_TYPE_IQ4_XS) { - return true; - } + ++compatible_consumer_count; + total_output_size += candidate_match.output_size; } - return false; + return compatible_consumer_count >= 2 && total_output_size >= std::max(match.input_size, 4096); +} + +static bool accepts_factored_tiled_f16_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, &consumer, kMulMatTiledF32F32Kernel, false); + return match.matched() && match.input->id == input.id && common_mul_mat_is_tiled_route(match); } static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) { @@ -391,20 +514,14 @@ static bool is_qwen_endpoint_rmsnorm_projection(const Graph & graph, const Commo } static bool match_symmetric_i4_low_row_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - // Disabled: this symmetric matmul route causes a substantial numeric performance regression. - (void) context; - (void) dispatch_match; - return false; - CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSymmetricI4LowRowWmmaKernel, false); if (!match.matched()) { match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSymmetricI4LowRowWmmaKernel, true); } - if (!match.matched() || (match.weight->type != GGML_TYPE_Q5_K && match.weight->type != GGML_TYPE_IQ4_XS) || - match.weight->alias_source.value >= 0 || match.token_count < 1 || match.token_count > 16 || - match.input_size % 64 != 0 || match.output_size % 64 != 0) { + if (!match.matched() || + !common_symmetric_i4_lowrow_mul_mat_eligible(context.graph, *context.root_node, *match.input)) { return false; } const CommonSymmetricI4ActivationLayout activation_layout = @@ -1140,16 +1257,54 @@ static bool build_mul_mat_dispatch( const CommonMulMatMatch & match, DispatchMatch & dispatch_match, CommonQ8ActivationPolicy q8_policy = CommonQ8ActivationPolicy::ExistingAlternateOnly) { - const bool publish_f16_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && - common_mul_mat_has_f16_alternate_consumer(context.graph, match); - const bool publish_q8_alternate = match.kernel.id == kMulMatTiledF32F32Kernel.id && - common_mul_mat_has_q8_alternate_consumer(context.graph, match); + const CommonActivationPublicationCapabilities capabilities = { + match.kernel.id == kMulMatTiledF32F32Kernel.id ? + DispatchActivationInputF16Row | DispatchActivationInputQ8_1X4 : + DispatchActivationInputNone, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + const CommonActivationPublicationDemand demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + const CommonActivationPublicationDemand f16_demand = + demand.format == CommonActivationPublicationFormat::F16Row ? demand : CommonActivationPublicationDemand{}; + const CommonActivationPublicationDemand q8_demand = + demand.format == CommonActivationPublicationFormat::Q8_1X4 ? demand : CommonActivationPublicationDemand{}; + const bool publish_f16_alternate = f16_demand.matched(); + const bool publish_q8_alternate = q8_demand.matched(); + const IQCodebookResourceSpec * codebook = iq_codebook_for_weight_format(match.weight_format); + if (codebook != nullptr && iq_codebook_matmul_required(match.weight_format) && + iq_codebook_matmul_disabled(match.weight_format)) { + return false; + } + const bool use_iq_codebook = + codebook != nullptr && !publish_f16_alternate && !iq_codebook_matmul_disabled(match.weight_format); + if ((match.weight_format == CommonMulMatWeightFormat::IQ1_S || + match.weight_format == CommonMulMatWeightFormat::IQ1_M || + match.weight_format == CommonMulMatWeightFormat::IQ3_XXS) && + (!use_iq_codebook || common_mul_mat_is_skinny_route(match))) { + return false; + } + const bool use_aligned_dynamic = match.token_count % 32 == 0 && + match.kernel.id == kMulMatTiledF32F32Kernel.id && !publish_f16_alternate && + !use_iq_codebook; + 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; Dispatch dispatch; - dispatch.kernel = - make_kernel_specialization(publish_f16_alternate ? kMulMatTiledF32F16AlternateKernel : match.kernel); + 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); + if (use_aligned_dynamic) { + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + } dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", - common_to_config_value(match.token_count)); + if (!use_aligned_dynamic) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + } dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.input_size", common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_size", common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_accumulation", "0"); @@ -1194,20 +1349,23 @@ static bool build_mul_mat_dispatch( } else { dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); } + if (use_iq_codebook) { + const ValueId grid = require_iq_codebook(context, dispatch_match, *codebook); + if (grid.value < 0) { + return false; + } + dispatch.bindings.push_back({ grid, 0, codebook->data.size() }); + } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (publish_f16_alternate) { - const size_t bytes = match.output->byte_count / 2; - const ValueId f16_output(context.next_plan_value.value + - static_cast(dispatch_match.transients.size())); constexpr const char * name = "common.mul_mat.tiled.f16"; - dispatch_match.transients.push_back({ f16_output, name, bytes, 256 }); - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { match.output->id, f16_output, GGML_TYPE_F16, bytes, name }, status)) { - dispatch_match.status.append(status); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, f16_demand, name, + publication) || + !common_append_activation_publication(dispatch_match, publication, name)) { return false; } - dispatch.bindings.push_back({ f16_output, 0, bytes }); + dispatch.bindings.push_back(publication.binding()); } dispatch_match.covered_nodes.push_back(context.root_index); @@ -1216,10 +1374,13 @@ static bool build_mul_mat_dispatch( } dispatch_match.dispatches.push_back(std::move(dispatch)); if (publish_q8_alternate) { - const size_t bytes = static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.output_size); - const ValueId q8_output(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); constexpr const char * name = "common.mul_mat.tiled.q8_1_x4"; - dispatch_match.transients.push_back({ q8_output, name, bytes, 256 }); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, q8_demand, name, + publication) || + !common_append_activation_publication(dispatch_match, publication, name)) { + return false; + } Dispatch quantize; quantize.kernel = make_kernel_specialization(kQuantizeQ8_1X4F32Kernel); @@ -1228,14 +1389,7 @@ static bool build_mul_mat_dispatch( quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", common_to_config_value(match.token_count * match.output_size / 128)); quantize.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - quantize.bindings.push_back({ q8_output, 0, bytes }); - - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { match.output->id, q8_output, GGML_TYPE_Q8_1, bytes, name }, status)) { - dispatch_match.status.append(status); - return false; - } + quantize.bindings.push_back(publication.binding()); dispatch_match.dispatches.push_back(std::move(quantize)); } return true; @@ -1307,108 +1461,30 @@ static bool vector_postops_route_supports(const MulMatPostOpsMatch & match) { match.input_size % 32 == 0 && match.output_size <= vector_max_output_size(match.weight_format); } -static size_t vector_publish_byte_count(ggml_type type, int64_t token_count, int64_t output_size) { - if (token_count <= 0 || output_size <= 0) { - return 0; - } - return static_cast(token_count) * ggml_row_size(type, output_size); -} - -static bool is_vector_packed_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { - if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || - consumer->inputs[1] != input.id) { - return false; - } - - const Value * weight = common_graph_value(graph, consumer->inputs[0]); - const Value * output = common_graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || - !input.contiguous || !weight->contiguous || !output->contiguous || weight->alias_source.value >= 0) { - return false; - } - - const bool supported_weight = weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_Q6_K; - return supported_weight && input.ne[1] >= 1 && input.ne[1] <= 5 && input.ne[0] >= 256 && input.ne[0] <= 32768 && - input.ne[0] % 256 == 0 && input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[0] == input.ne[0] && - weight->ne[1] >= 64 && weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && - weight->ne[3] == 1 && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && - output->ne[3] == 1; -} - -static bool has_vector_packed_q8_consumer(const Graph & graph, const Value & value) { - if (!graph.has_index()) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(value.id)) { - if (is_vector_packed_q8_consumer(graph, consumer, value)) { - return true; - } - } - return false; -} - -static bool is_vector_f16_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { - if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || - consumer->inputs[1] != input.id) { - return false; - } - - const Value * weight = common_graph_value(graph, consumer->inputs[0]); - const Value * output = common_graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || - !input.contiguous || !weight->contiguous || !output->contiguous || weight->alias_source.value >= 0) { - return false; - } - - CommonMulMatWeightFormat format; - if (!common_mul_mat_format_for_type(weight->type, format)) { - return false; - } - return input.ne[0] == weight->ne[0] && output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && - input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1 && output->ne[2] == 1 && - output->ne[3] == 1 && common_mul_mat_uses_k16_major_f16(format, input.ne[0], output->ne[0], input.ne[1]); -} - -static bool has_vector_f16_consumer(const Graph & graph, const Value & value) { - if (!graph.has_index()) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(value.id)) { - if (is_vector_f16_consumer(graph, consumer, value)) { - return true; - } +static bool prepare_vector_publish_plan(const DispatchMatchContext & context, + const Value & output, + DispatchMatch & dispatch_match, + VectorPublishPlan & plan) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4 | DispatchActivationInputF16Row, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, output, capabilities); + plan.demand = publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (!plan.demand.matched()) { + return true; } - return false; -} -static VectorPublishPlan make_vector_publish_plan(const DispatchMatchContext & context, - const Value & output, - int64_t token_count, - int64_t output_size, - const DispatchMatch & dispatch_match) { - VectorPublishPlan plan; - if (has_vector_packed_q8_consumer(context.graph, output)) { - plan.publish_q8 = true; - plan.q8_output = - ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); - plan.q8_byte_count = vector_publish_byte_count(GGML_TYPE_Q8_1, token_count, output_size); - } else if (has_vector_f16_consumer(context.graph, output)) { - plan.fused_format = VectorPublishFormat::F16; - plan.fused_output = - ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); - plan.fused_byte_count = vector_publish_byte_count(GGML_TYPE_F16, token_count, output_size); - } - return plan; + const char * name = plan.demand.format == CommonActivationPublicationFormat::Q8_1X4 ? + "common.mul_mat.vector.q8_1_x4" : + "common.mul_mat.vector.f16"; + return common_reserve_activation_publication(context, dispatch_match, output, plan.demand, name, + plan.publication) && + common_append_activation_publication(dispatch_match, plan.publication, name); } -static int64_t vector_publish_format_config_value(VectorPublishFormat format) { - switch (format) { - case VectorPublishFormat::F16: - return GGML_TYPE_F16; - case VectorPublishFormat::None: - return 0; - } - return 0; +static int64_t vector_publish_format_config_value(const VectorPublishPlan & plan) { + return plan.publication.format == CommonActivationPublicationFormat::F16Row ? GGML_TYPE_F16 : 0; } static bool append_vector_q8_publish(const Value & output, @@ -1416,14 +1492,10 @@ static bool append_vector_q8_publish(const Value & output, int64_t output_size, VectorPublishPlan & publish, DispatchMatch & dispatch_match) { - if (!publish.publish_q8) { + if (publish.publication.format != CommonActivationPublicationFormat::Q8_1X4) { return true; } - if (publish.q8_byte_count == 0) { - return false; - } - constexpr const char * name = "common.mul_mat.vector.q8_1_x4"; Dispatch quantize; quantize.kernel = make_kernel_specialization(kQuantizeQ8_1X4F32Kernel); quantize.kernel.integer_parameters.emplace("token_count", token_count); @@ -1431,43 +1503,14 @@ static bool append_vector_q8_publish(const Value & output, quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", common_to_config_value(token_count * output_size / 128)); quantize.bindings.push_back({ output.id, 0, output.byte_count }); - quantize.bindings.push_back({ publish.q8_output, 0, publish.q8_byte_count }); + quantize.bindings.push_back(publish.publication.binding()); dispatch_match.dispatches.push_back(std::move(quantize)); - dispatch_match.transients.push_back({ publish.q8_output, name, publish.q8_byte_count, 256 }); - - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { output.id, publish.q8_output, GGML_TYPE_Q8_1, publish.q8_byte_count, name }, status)) { - dispatch_match.status.append(status); - return false; - } - return true; -} - -static bool append_vector_f16_publish_metadata(const Value & output, - VectorPublishPlan & publish, - DispatchMatch & dispatch_match) { - if (publish.fused_format != VectorPublishFormat::F16) { - return true; - } - if (publish.fused_byte_count == 0) { - return false; - } - - constexpr const char * name = "common.mul_mat.vector.f16"; - dispatch_match.transients.push_back({ publish.fused_output, name, publish.fused_byte_count, 256 }); - Status status; - if (!dispatch_match.metadata.append_alternate_value( - { output.id, publish.fused_output, GGML_TYPE_F16, publish.fused_byte_count, name }, status)) { - dispatch_match.status.append(status); - return false; - } return true; } static DispatchBinding vector_next_output_binding(const Value & output, const VectorPublishPlan & publish) { - if (publish.fused_format == VectorPublishFormat::F16) { - return { publish.fused_output, 0, publish.fused_byte_count }; + if (publish.publication.format == CommonActivationPublicationFormat::F16Row) { + return publish.publication.binding(); } return { output.id, 0, output.byte_count }; } @@ -1485,25 +1528,37 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kMulMatVectorF32F32Kernel); - dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", - common_to_config_value(match.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", - common_to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", - common_to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); - dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_unary_op", - std::to_string(unary_kind_config_value(match.output_unary_op))); - VectorPublishPlan publish = - make_vector_publish_plan(context, *match.output, match.token_count, match.output_size, dispatch_match); - if (!append_vector_f16_publish_metadata(*match.output, publish, dispatch_match)) { + const IQCodebookResourceSpec * codebook = iq_codebook_for_weight_format(match.weight_format); + if (codebook != nullptr && iq_codebook_matmul_required(match.weight_format) && + iq_codebook_matmul_disabled(match.weight_format)) { + return false; + } + VectorPublishPlan publish; + if (!prepare_vector_publish_plan(context, *match.output, dispatch_match, publish)) { return false; } - dispatch.kernel.compile_parameters.emplace( - "ggml.matmul.vector.publish_format", - common_to_config_value(vector_publish_format_config_value(publish.fused_format))); + + if (match.weight_format == CommonMulMatWeightFormat::Q1_0 && match.weight->alias_source.value < 0 && + match.input_size % 128 == 0 && !match.has_fused_unary && !publish.demand.matched()) { + DispatchBinding q8_activation; + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + q8_activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + dispatch.kernel = make_kernel_specialization( + match.output_size <= 1024 ? kMulMatVectorQ1_0Q8_1X4F32Kernel : + kMulMatVectorQ1_0Q8_1X4Rows2F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.q1_0_q8_1_x4.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.q1_0_q8_1_x4.output_size", + common_to_config_value(match.output_size)); + dispatch.bindings.push_back(q8_activation); + 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; + } + } DispatchBinding activation; CommonMulMatWeightFormat weight_format = match.weight_format; @@ -1515,6 +1570,31 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, return false; } + const bool use_iq_codebook = + codebook != nullptr && !iq_codebook_matmul_disabled(match.weight_format) && !uses_q8_activation && + publish.publication.format != CommonActivationPublicationFormat::F16Row; + if ((match.weight_format == CommonMulMatWeightFormat::IQ1_S || + match.weight_format == CommonMulMatWeightFormat::IQ1_M || + match.weight_format == CommonMulMatWeightFormat::IQ3_XXS) && + !use_iq_codebook) { + return false; + } + dispatch.kernel = make_kernel_specialization( + use_iq_codebook ? vector_iq_codebook_kernel(match.weight_format) : kMulMatVectorF32F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.output_unary_op", + std::to_string(unary_kind_config_value(match.output_unary_op))); + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector.publish_format", + common_to_config_value(vector_publish_format_config_value(publish))); + dispatch.kernel.compile_parameters.emplace( "ggml.matmul.vector.weight_format", common_to_config_value(common_mul_mat_format_config_value(weight_format))); if (uses_q8_activation) { @@ -1529,6 +1609,13 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, } else { dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); } + if (use_iq_codebook) { + const ValueId grid = require_iq_codebook(context, dispatch_match, *codebook); + if (grid.value < 0) { + return false; + } + dispatch.bindings.push_back({ grid, 0, codebook->data.size() }); + } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch.bindings.push_back(vector_next_output_binding(*match.output, publish)); @@ -1543,6 +1630,9 @@ static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & context, const MulMatPostOpsMatch & match, DispatchMatch & dispatch_match) { + if (iq_codebook_matmul_required(match.weight_format)) { + return false; + } if (match.token_count != 1 || match.input_size % 32 != 0 || match.output_size > vector_max_output_size(match.weight_format)) { return false; @@ -1571,14 +1661,13 @@ static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & c dispatch.kernel.compile_parameters.emplace( "ggml.matmul.vector_postops.apply_residual", common_to_config_value(match.has_residual ? int64_t{ 1 } : int64_t{ 0 })); - VectorPublishPlan publish = - make_vector_publish_plan(context, *match.residual_output, match.token_count, match.output_size, dispatch_match); - if (!append_vector_f16_publish_metadata(*match.residual_output, publish, dispatch_match)) { + VectorPublishPlan publish; + if (!prepare_vector_publish_plan(context, *match.residual_output, dispatch_match, publish)) { return false; } dispatch.kernel.compile_parameters.emplace( "ggml.matmul.vector.publish_format", - common_to_config_value(vector_publish_format_config_value(publish.fused_format))); + common_to_config_value(vector_publish_format_config_value(publish))); DispatchBinding activation; CommonMulMatWeightFormat weight_format = match.weight_format; @@ -1693,15 +1782,26 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, if (!match.matched()) { return false; } + if (iq_codebook_matmul_required(match.weight_format)) { + return false; + } if (vector_postops_route_supports(match)) { return false; } + const bool use_aligned_dynamic = match.token_count % 32 == 0 && + match.kernel.id == kMulMatTiledBiasAddF32F32Kernel.id; Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel = make_kernel_specialization( + use_aligned_dynamic ? kMulMatTiledAlignedBiasAddF32F32Kernel : match.kernel); + if (use_aligned_dynamic) { + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + } dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", - common_to_config_value(match.token_count)); + if (!use_aligned_dynamic) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + } dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.input_size", common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.output_size", @@ -1941,6 +2041,9 @@ static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, if (!match.matched()) { return false; } + if (iq_codebook_matmul_required(match.weight_format)) { + return false; + } if (vector_mul_mat_route_supports(match)) { return false; } @@ -1965,6 +2068,9 @@ static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & conte if (!match.matched()) { return false; } + if (iq_codebook_matmul_required(match.root.weight_format)) { + return false; + } MulMatPostOpsMatch postops; postops.input = match.root.input; postops.weight = match.root.weight; @@ -1984,6 +2090,89 @@ static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & conte } void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { + registry.add_activation_consumer({ + "common.mul_mat.input.f16_row.factored_tiled", + GGML_OP_MUL_MAT, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Tiled, + DispatchActivationProducerRequirementFactoredTiledF16, + DispatchSource::Common, + accepts_factored_tiled_f16_input, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.q8_1_x4.q1_0_preferred_publication", + GGML_OP_MUL_MAT, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + accepts_preferred_q1_0_q8_publication, + DispatchActivationConsumerTraitQ8PublicationPreferred, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.q8_1_x4.preferred_publication", + GGML_OP_MUL_MAT, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + accepts_preferred_q8_publication, + DispatchActivationConsumerTraitQ8PublicationPreferred, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.q8_1_x4.decode", + GGML_OP_MUL_MAT, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + common_accepts_q8_1_x4_decode_mul_mat, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.q8_1_x4.prefill", + GGML_OP_MUL_MAT, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + accepts_tiled_q8_publication, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.f16_k16_major", + GGML_OP_MUL_MAT, + DispatchActivationInputF16K16Major, + DispatchActivationConsumerUse::Tiled, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + common_accepts_f16_k16_major_mul_mat, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.f16_k16_major.packed_producer", + GGML_OP_MUL_MAT, + DispatchActivationInputF16K16Major, + DispatchActivationConsumerUse::Tiled, + DispatchActivationProducerRequirementPackedK16Economy, + DispatchSource::Common, + common_accepts_f16_k16_major_packed_producer_mul_mat, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.f16_row.k16_preparation", + GGML_OP_MUL_MAT, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Tiled, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + common_accepts_f16_k16_major_mul_mat, + }); + registry.add_activation_consumer({ + "common.mul_mat.input.f16_row.prefill", + GGML_OP_MUL_MAT, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Row, + DispatchActivationProducerRequirementNone, + DispatchSource::Common, + common_accepts_f16_row_prefill_mul_mat, + }); registry.add({ "common.mul_mat.vector_q6_selected_outputs", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp index 2a078613af33..b6d35c974bac 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp @@ -1,6 +1,7 @@ #include "dispatch-rmsnorm.h" #include "dispatch-binary-common.h" +#include "dispatch-activation-publication.h" #include "dispatch-layout-utils.h" #include "dispatch-mul-mat-common.h" #include "ggml.h" @@ -10,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -27,8 +29,8 @@ static constexpr KernelCatalogRef kRmsNormBinaryF32F16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32_f16"); static constexpr KernelCatalogRef kRmsNormBinaryF32K16Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32_k16"); -static constexpr KernelCatalogRef kRmsNormBinaryQ8_1X4F16Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_q8_1_x4_f16"); +static constexpr KernelCatalogRef kRmsNormBinaryQ8_1X4PublishKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_q8_1_x4_publish"); static constexpr KernelCatalogRef kRmsNormF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_f32"); static constexpr KernelCatalogRef kRmsNormMulRopeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_mul_rope_f32"); @@ -38,10 +40,8 @@ static constexpr KernelCatalogRef kRmsNormBinarySymmetricI4K32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_symmetric_i4_k32"); static constexpr KernelCatalogRef kRmsNormGateSiluMulSymmetricI4K32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32"); -static constexpr KernelCatalogRef kRmsNormGateF32F16Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_f16"); -static constexpr KernelCatalogRef kRmsNormGateF32Q8_1X4Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_q8_1_x4"); +static constexpr KernelCatalogRef kRmsNormGateF32PublishKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_publish"); static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); @@ -98,47 +98,6 @@ static bool supported_rmsnorm_input_layout(const Value & input, return strided_f32_storage_span_bytes(input, input_span_bytes); } -static bool is_packed_q8_consumer(const Graph & graph, const GraphNode * consumer, const Value & input) { - if (consumer == nullptr || consumer->op != GGML_OP_MUL_MAT || consumer->inputs.size() != 2 || - consumer->inputs[1] != input.id) { - return false; - } - const Value * weight = graph_value(graph, consumer->inputs[0]); - const Value * output = graph_value(graph, consumer->output); - if (weight == nullptr || output == nullptr || input.type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || - !input.contiguous || !weight->contiguous || !output->contiguous) { - return false; - } - const bool prefill = (weight->type == GGML_TYPE_Q5_K || weight->type == GGML_TYPE_IQ4_XS) && input.ne[1] >= 256 && - input.ne[1] <= 2048 && input.ne[1] % 256 == 0; - const bool decode = input.ne[1] >= 1 && input.ne[1] <= 5 && weight->alias_source.value < 0 && - (weight->type == GGML_TYPE_Q4_K || - (weight->type == GGML_TYPE_Q6_K && !graph.index().consumers(output->id).empty())); - return (prefill || decode) && input.ne[0] >= 256 && input.ne[0] <= 32768 && input.ne[0] % 256 == 0 && - input.ne[2] == 1 && input.ne[3] == 1 && weight->ne[0] == input.ne[0] && weight->ne[1] >= 64 && - weight->ne[1] <= 262144 && weight->ne[1] % 64 == 0 && weight->ne[2] == 1 && weight->ne[3] == 1 && - output->ne[0] == weight->ne[1] && output->ne[1] == input.ne[1] && output->ne[2] == 1 && output->ne[3] == 1; -} - -static bool has_packed_q8_consumer(const Graph & graph, const Value & value) { - if (!graph.has_index()) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(value.id)) { - if (is_packed_q8_consumer(graph, consumer, value)) { - return true; - } - } - return false; -} - -static size_t q8_1_x4_byte_count(int64_t token_count, int64_t hidden_size) { - if (token_count <= 0 || hidden_size <= 0) { - return 0; - } - return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); -} - static bool is_weight_shape(const Value & weight, int64_t hidden_size) { if (weight.ne[0] != hidden_size) { return false; @@ -599,7 +558,6 @@ static AddRmsNormBinarySymmetricI4Match match_add_rmsnorm_binary_symmetric_i4(co RmsNormBinaryMatch rms_binary = match_rmsnorm_binary_f32(graph, rms_node, rms_node_index); if (!rms_binary.matched() || rms_binary.input->id != residual->id || rms_binary.op != BinaryKind::Mul || !is_supported_symmetric_i4_hidden_size(rms_binary.hidden_size) || - !common_has_symmetric_i4_lowrow_consumer(graph, *rms_binary.output) || !pairwise_distinct_storage_roots( std::array{ lhs, rhs, residual, rms_binary.rhs, rms_binary.output })) { return {}; @@ -757,12 +715,22 @@ static bool match_add_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchCo return false; } - const CommonSymmetricI4ActivationLayout activation_layout = - common_symmetric_i4_activation_layout(fused.rms_binary.hidden_size, fused.rms_binary.token_count); - if (activation_layout.total_bytes == 0) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputSymmetricI4K32, + DispatchActivationInputSymmetricI4K32, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *fused.rms_binary.output, capabilities); + if (!publication_plan.matched()) { + return false; + } + const CommonActivationPublicationDemand & demand = publication_plan.preferred(); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *fused.rms_binary.output, demand, + kCommonSymmetricI4K32ActivationAlternateName, publication) || + !common_append_activation_publication(match, publication, kCommonSymmetricI4K32ActivationAlternateName)) { return false; } - const ValueId activation = context.next_plan_value; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kAddRmsNormBinarySymmetricI4K32Kernel); @@ -776,25 +744,14 @@ static bool match_add_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchCo dispatch.bindings.push_back({ fused.residual->id, 0, fused.residual->byte_count }); dispatch.bindings.push_back({ fused.rms_binary.rhs->id, 0, fused.rms_binary.rhs->byte_count }); dispatch.bindings.push_back({ fused.rms_binary.output->id, 0, fused.rms_binary.output->byte_count }); - dispatch.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); - - Status metadata_status; - if (!match.metadata.append_alternate_value( - { fused.rms_binary.output->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, - kCommonSymmetricI4K32ActivationAlternateName }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(publication.binding(0, publication.payload_bytes)); + dispatch.bindings.push_back(publication.binding(publication.scales_offset, publication.scales_bytes)); + dispatch.bindings.push_back(publication.binding(publication.sums_offset, publication.sums_bytes)); match.covered_nodes.push_back(fused.add_node_index); match.covered_nodes.push_back(fused.rms_binary.rms_node_index); match.covered_nodes.push_back(fused.rms_binary.binary_node_index); match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back( - { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); return match.status.success(); } @@ -809,17 +766,26 @@ static bool match_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchContex fused.token_count > 16 || fused.rms_node_index >= context.covered_nodes.size() || fused.binary_node_index >= context.covered_nodes.size() || context.covered_nodes[fused.rms_node_index] || context.covered_nodes[fused.binary_node_index] || - !common_has_symmetric_i4_lowrow_consumer(context.graph, *fused.output) || !pairwise_distinct_storage_roots(std::array{ fused.input, fused.rhs, fused.output })) { return false; } - const CommonSymmetricI4ActivationLayout activation_layout = - common_symmetric_i4_activation_layout(fused.hidden_size, fused.token_count); - if (activation_layout.total_bytes == 0) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputSymmetricI4K32, + DispatchActivationInputSymmetricI4K32, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *fused.output, capabilities); + if (!publication_plan.matched()) { + return false; + } + const CommonActivationPublicationDemand & demand = publication_plan.preferred(); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *fused.output, demand, + kCommonSymmetricI4K32ActivationAlternateName, publication) || + !common_append_activation_publication(match, publication, kCommonSymmetricI4K32ActivationAlternateName)) { return false; } - const ValueId activation = context.next_plan_value; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRmsNormBinarySymmetricI4K32Kernel); @@ -831,24 +797,13 @@ static bool match_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchContex dispatch.bindings.push_back({ fused.input->id, 0, fused.input->byte_count }); dispatch.bindings.push_back({ fused.rhs->id, 0, fused.rhs->byte_count }); dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); - dispatch.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); - - Status metadata_status; - if (!match.metadata.append_alternate_value( - { fused.output->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, - kCommonSymmetricI4K32ActivationAlternateName }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(publication.binding(0, publication.payload_bytes)); + dispatch.bindings.push_back(publication.binding(publication.scales_offset, publication.scales_bytes)); + dispatch.bindings.push_back(publication.binding(publication.sums_offset, publication.sums_bytes)); match.covered_nodes.push_back(fused.rms_node_index); match.covered_nodes.push_back(fused.binary_node_index); match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back( - { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); return match.status.success(); } @@ -858,8 +813,13 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di return false; } - const bool use_i4 = - fused.gate_op == UnaryKind::Silu && common_has_symmetric_i4_lowrow_consumer(context.graph, *fused.output); + const CommonActivationPublicationCapabilities i4_capabilities = { + fused.gate_op == UnaryKind::Silu ? DispatchActivationInputSymmetricI4K32 : DispatchActivationInputNone, + DispatchActivationInputSymmetricI4K32, + }; + const CommonActivationPublicationPlan i4_plan = + common_select_activation_publication_plan(context, *fused.output, i4_capabilities); + const bool use_i4 = i4_plan.matched(); if (!use_i4) { if (fused.hidden_size > 1024) { return false; @@ -877,23 +837,48 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di consumer = common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_MUL_MAT); } } - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(context.graph, consumer, kRmsNormGateF32F16Kernel, false); - const bool packed_output = projection.matched() && projection.input->id == packed_input->id && - projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, - projection.output_size, projection.token_count); - const bool q8_output = - !packed_output && packed_input->ne[1] <= 5 && is_packed_q8_consumer(context.graph, consumer, *packed_input); - const size_t activation_bytes = - q8_output ? q8_1_x4_byte_count(fused.token_count, fused.hidden_size) : fused.output->byte_count / 2; - const ValueId activation = context.next_plan_value; + const uint32_t enabled_formats = DispatchActivationInputF16K16Major | DispatchActivationInputF16Row; + const CommonActivationPublicationCapabilities capabilities = { + enabled_formats, + DispatchActivationInputF16Row, + DispatchActivationProducerRequirementFactoredTiledF16, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *packed_input, capabilities); + const CommonActivationPublicationCapabilities q8_capabilities = { + packed_input->ne[1] <= 5 ? DispatchActivationInputQ8_1X4 : DispatchActivationInputNone, + DispatchActivationInputQ8_1X4, + DispatchActivationProducerRequirementNone, + DispatchActivationInputNone, + DispatchActivationConsumerSourceAll, + DispatchActivationConsumerTraitQ8PublicationPreferred, + }; + const CommonActivationPublicationPlan q8_plan = + common_select_activation_publication_plan(context, *packed_input, q8_capabilities); + const CommonActivationPublicationDemand * publication_demand = + q8_plan.matched() ? &q8_plan.preferred() : + publication_plan.matched() ? &publication_plan.preferred() : nullptr; + const bool packed_output = publication_demand != nullptr && + publication_demand->format == CommonActivationPublicationFormat::F16K16Major; + const bool q8_output = publication_demand != nullptr && + publication_demand->format == CommonActivationPublicationFormat::Q8_1X4; const char * activation_name = q8_output ? "common.rmsnorm_gate.q8_1_x4" : - packed_output ? "common.rmsnorm_gate.k16_major_f16" : - "common.rmsnorm_gate.f16"; + packed_output ? "common.rmsnorm_gate.k16_major_f16" : "common.rmsnorm_gate.f16"; + const Value & publication_source = packed_output ? *packed_input : *fused.output; + CommonActivationPublication publication; + if (publication_demand != nullptr) { + if (!common_reserve_activation_publication(context, match, publication_source, *publication_demand, + activation_name, publication) || + !common_append_activation_publication(match, publication, activation_name)) { + return false; + } + } else if (!common_reserve_private_activation_output(context, match, *fused.output, + CommonActivationPublicationFormat::F16Row, + activation_name, publication)) { + return false; + } Dispatch dispatch; - dispatch.kernel = - make_kernel_specialization(q8_output ? kRmsNormGateF32Q8_1X4Kernel : kRmsNormGateF32F16Kernel); + dispatch.kernel = make_kernel_specialization(kRmsNormGateF32PublishKernel); dispatch.kernel.integer_parameters.emplace("token_count", fused.token_count); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.hidden_size", to_config_value(fused.hidden_size)); @@ -901,42 +886,29 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.gate_op", std::to_string(unary_kind_config_value(fused.gate_op))); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.f16_output_row_width", - to_config_value(packed_output ? projection.input_size : 0)); + to_config_value(packed_output ? publication_demand->consumer_value->ne[0] : 0)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.publish_q8", q8_output ? "1" : "0"); dispatch.bindings.push_back({ fused.input->id, 0, fused.input->byte_count }); dispatch.bindings.push_back({ fused.weight->id, 0, fused.weight->byte_count }); dispatch.bindings.push_back({ fused.raw_gate->id, 0, fused.raw_gate->byte_count }); dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); - dispatch.bindings.push_back({ activation, 0, activation_bytes }); - Status status; - const bool recorded = - packed_output ? - match.metadata.append_generated_resource( - { packed_input->id, GeneratedResourceRole::F16K16Major, activation, activation_bytes, {} }, - status) : - match.metadata.append_alternate_value( - { fused.output->id, activation, q8_output ? GGML_TYPE_Q8_1 : GGML_TYPE_F16, activation_bytes, - activation_name }, - status); - if (!recorded) { - match.status.append(status); - return false; - } + dispatch.bindings.push_back(publication.binding()); for (const GraphNode * covered : fused.covered) { if (!append_covered_node_index_once(context.graph, context.covered_nodes, covered, match.covered_nodes)) { return false; } } - match.transients.push_back({ activation, activation_name, activation_bytes, 256 }); match.dispatches.push_back(std::move(dispatch)); return true; } - const CommonSymmetricI4ActivationLayout activation_layout = - common_symmetric_i4_activation_layout(fused.hidden_size, fused.token_count); - if (activation_layout.total_bytes == 0) { + const CommonActivationPublicationDemand & i4_demand = i4_plan.preferred(); + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *fused.output, i4_demand, + kCommonSymmetricI4K32ActivationAlternateName, publication) || + !common_append_activation_publication(match, publication, kCommonSymmetricI4K32ActivationAlternateName)) { return false; } - const ValueId activation = context.next_plan_value; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRmsNormGateSiluMulSymmetricI4K32Kernel); @@ -949,75 +921,31 @@ static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, Di dispatch.bindings.push_back({ fused.weight->id, 0, fused.weight->byte_count }); dispatch.bindings.push_back({ fused.raw_gate->id, 0, fused.raw_gate->byte_count }); dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); - dispatch.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); - dispatch.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); - - Status metadata_status; - if (!match.metadata.append_alternate_value( - { fused.output->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, - kCommonSymmetricI4K32ActivationAlternateName }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(publication.binding(0, publication.payload_bytes)); + dispatch.bindings.push_back(publication.binding(publication.scales_offset, publication.scales_bytes)); + dispatch.bindings.push_back(publication.binding(publication.sums_offset, publication.sums_bytes)); for (const GraphNode * covered : fused.covered) { if (!append_covered_node_index_once(context.graph, context.covered_nodes, covered, match.covered_nodes)) { return false; } } match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back( - { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); return match.status.success(); } -static bool has_qualified_rmsnorm_k16_consumer(const Graph & graph, const RmsNormBinaryMatch & rms) { +static bool rmsnorm_k16_publication_enabled(const RmsNormBinaryMatch & rms) { const bool qualified_shape = - (rms.token_count == 512 && (rms.hidden_size == 2048 || rms.hidden_size == 4096 || rms.hidden_size == 5120 || - rms.hidden_size == 6144 || rms.hidden_size == 8192)) || + (rms.token_count == 512 && (rms.hidden_size == 2048 || rms.hidden_size == 3072 || rms.hidden_size == 4096 || + rms.hidden_size == 5120 || rms.hidden_size == 6144 || rms.hidden_size == 8192)) || (rms.token_count == 1024 && rms.hidden_size == 5120); - if (!qualified_shape || rms.op != BinaryKind::Mul || rms.output->ne[1] != rms.token_count || - rms.output->ne[2] != 1 || rms.output->ne[3] != 1 || - !pairwise_distinct_storage_roots(std::array{ rms.input, rms.rhs, rms.output })) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(rms.output->id)) { - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(graph, consumer, kRmsNormBinaryF32K16Kernel, false); - if (projection.matched() && projection.input->id == rms.output->id && - projection.weight->alias_source.value < 0 && - common_mul_mat_uses_k16_major_f16(projection.weight_format, projection.input_size, projection.output_size, - projection.token_count)) { - return true; - } - } - return false; + return qualified_shape && rms.op == BinaryKind::Mul && rms.output->ne[1] == rms.token_count && + rms.output->ne[2] == 1 && rms.output->ne[3] == 1 && + pairwise_distinct_storage_roots(std::array{ rms.input, rms.rhs, rms.output }); } -static bool has_qualified_rmsnorm_f16_consumer(const Graph & graph, const RmsNormBinaryMatch & rms) { - if (!graph.has_index() || rms.token_count < 128 || rms.output->ne[1] != rms.token_count || rms.output->ne[2] != 1 || - rms.output->ne[3] != 1) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(rms.output->id)) { - const CommonMulMatMatch projection = - common_match_mul_mat_any_format(graph, consumer, kRmsNormBinaryF32F16Kernel, false); - if (!projection.matched() || projection.input->id != rms.output->id || - projection.input_size != rms.hidden_size || projection.token_count != rms.token_count || - projection.input_size % 256 != 0 || projection.output_size % 64 != 0 || - projection.weight->alias_source.value >= 0) { - continue; - } - if (projection.weight->type == GGML_TYPE_Q6_K && projection.token_count % 128 == 0) { - return true; - } - if (projection.weight->type == GGML_TYPE_Q4_K && projection.token_count % 256 == 0 && - projection.output_size >= projection.input_size / 4) { - return true; - } - } - return false; +static bool rmsnorm_f16_publication_enabled(const RmsNormBinaryMatch & rms) { + return rms.token_count >= 128 && rms.output->ne[1] == rms.token_count && rms.output->ne[2] == 1 && + rms.output->ne[3] == 1; } static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { @@ -1034,9 +962,18 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte } const bool strided_input = rms_match.input_stride != rms_match.hidden_size; - const bool packed_output = !strided_input && has_qualified_rmsnorm_k16_consumer(context.graph, rms_match); - const bool f16_output = - !strided_input && !packed_output && has_qualified_rmsnorm_f16_consumer(context.graph, rms_match); + const uint32_t enabled_formats = + (!strided_input && rmsnorm_k16_publication_enabled(rms_match) ? DispatchActivationInputF16K16Major : 0) | + (!strided_input && rmsnorm_f16_publication_enabled(rms_match) ? DispatchActivationInputF16Row : 0); + const CommonActivationPublicationCapabilities capabilities = { enabled_formats }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *rms_match.output, capabilities); + const CommonActivationPublicationDemand * publication_demand = + publication_plan.matched() ? &publication_plan.preferred() : nullptr; + const bool packed_output = publication_demand != nullptr && + publication_demand->format == CommonActivationPublicationFormat::F16K16Major; + const bool f16_output = publication_demand != nullptr && + publication_demand->format == CommonActivationPublicationFormat::F16Row; Dispatch dispatch; dispatch.kernel = make_kernel_specialization(strided_input ? kRmsNormBinaryStridedF32Kernel : packed_output ? kRmsNormBinaryF32K16Kernel : @@ -1060,34 +997,22 @@ static bool match_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & conte dispatch.bindings.push_back({ rms_match.rhs->id, 0, rms_match.rhs->byte_count }); dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); - if (packed_output) { - const ValueId packed = context.next_plan_value; - const size_t bytes = rms_match.output->byte_count / 2; - Status status; - if (!match.metadata.append_generated_resource( - { rms_match.output->id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, status)) { - match.status.append(status); - return false; - } - dispatch.bindings.push_back({ packed, 0, bytes }); - match.transients.push_back({ packed, "common.rmsnorm_binary.k16_major_f16", bytes, 256 }); - } else if (f16_output) { - const ValueId f16 = context.next_plan_value; - const size_t bytes = rms_match.output->byte_count / 2; - Status status; - if (!match.metadata.append_alternate_value( - { rms_match.output->id, f16, GGML_TYPE_F16, bytes, "common.rmsnorm_binary.f16" }, status)) { - match.status.append(status); + if (publication_demand != nullptr) { + const char * publication_name = packed_output ? "common.rmsnorm_binary.k16_major_f16" : + "common.rmsnorm_binary.f16"; + CommonActivationPublication publication; + if (!common_reserve_activation_publication(context, match, *rms_match.output, *publication_demand, + publication_name, publication) || + !common_append_activation_publication(match, publication, publication_name)) { return false; } - dispatch.bindings.push_back({ f16, 0, bytes }); - match.transients.push_back({ f16, "common.rmsnorm_binary.f16", bytes, 256 }); + dispatch.bindings.push_back(publication.binding()); } match.covered_nodes.push_back(rms_match.rms_node_index); match.covered_nodes.push_back(rms_match.binary_node_index); match.dispatches.push_back(std::move(dispatch)); - return true; + return match.status.success(); } static bool match_rmsnorm_mul_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { @@ -1125,24 +1050,43 @@ static bool match_rmsnorm_binary_q8_1_x4_dispatch(const DispatchMatchContext & c } const RmsNormBinaryMatch rms_match = match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index); + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + DispatchActivationProducerRequirementNone, + DispatchActivationInputQ8_1X4 | DispatchActivationInputF16Row, + DispatchActivationConsumerSourceAll, + DispatchActivationConsumerTraitQ8PublicationPreferred, + }; + const CommonActivationPublicationPlan publication_plan = + rms_match.matched() ? common_select_activation_publication_plan(context, *rms_match.output, capabilities) : + CommonActivationPublicationPlan{}; + const CommonActivationPublicationDemand * q8_demand = + publication_plan.find(CommonActivationPublicationFormat::Q8_1X4); + const CommonActivationPublicationDemand * f16_demand = + publication_plan.find(CommonActivationPublicationFormat::F16Row); if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || rms_match.binary_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || context.covered_nodes[rms_match.binary_node_index] || - !has_packed_q8_consumer(context.graph, *rms_match.output)) { + q8_demand == nullptr || f16_demand == nullptr) { return false; } - const size_t q8_byte_count = q8_1_x4_byte_count(rms_match.token_count, rms_match.hidden_size); - if (q8_byte_count == 0) { + CommonActivationPublication q8_publication; + if (!common_reserve_activation_publication(context, match, *rms_match.output, *q8_demand, + "common.rmsnorm_binary.q8_1_x4", q8_publication)) { + return false; + } + CommonActivationPublication f16_publication; + if (!common_reserve_activation_publication(context, match, *rms_match.output, *f16_demand, + "common.rmsnorm_binary.f16", f16_publication) || + !common_append_activation_publication(match, q8_publication, "common.rmsnorm_binary.q8_1_x4") || + !common_append_activation_publication(match, f16_publication, "common.rmsnorm_binary.f16")) { return false; } - - const ValueId q8_value = context.next_plan_value; - const ValueId f16_value(context.next_plan_value.value + 1); - const size_t f16_byte_count = rms_match.output->byte_count / 2; Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kRmsNormBinaryQ8_1X4F16Kernel); + dispatch.kernel = make_kernel_specialization(kRmsNormBinaryQ8_1X4PublishKernel); dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_q8_1_x4.hidden_size", to_config_value(rms_match.hidden_size)); @@ -1150,28 +1094,16 @@ static bool match_rmsnorm_binary_q8_1_x4_dispatch(const DispatchMatchContext & c to_config_value(rms_match.epsilon)); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_q8_1_x4.op", std::to_string(binary_kind_config_value(rms_match.op))); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); dispatch.bindings.push_back({ rms_match.rhs->id, 0, rms_match.rhs->byte_count }); dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); - dispatch.bindings.push_back({ q8_value, 0, q8_byte_count }); - dispatch.bindings.push_back({ f16_value, 0, f16_byte_count }); - - Status metadata_status; - if (!match.metadata.append_alternate_value( - { rms_match.output->id, q8_value, GGML_TYPE_Q8_1, q8_byte_count, "common.rmsnorm_binary.q8_1_x4" }, - metadata_status) || - !match.metadata.append_alternate_value( - { rms_match.output->id, f16_value, GGML_TYPE_F16, f16_byte_count, "common.rmsnorm_binary.f16" }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(q8_publication.binding()); + dispatch.bindings.push_back(f16_publication.binding()); match.covered_nodes.push_back(rms_match.rms_node_index); match.covered_nodes.push_back(rms_match.binary_node_index); match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back({ q8_value, "common.rmsnorm_binary.q8_1_x4", q8_byte_count, 256 }); - match.transients.push_back({ f16_value, "common.rmsnorm_binary.f16", f16_byte_count, 256 }); return match.status.success(); } @@ -1188,6 +1120,7 @@ static bool match_rmsnorm_f32_dispatch(const DispatchMatchContext & context, Dis Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRmsNormF32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_f32.hidden_size", to_config_value(rms_match.hidden_size)); dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_f32.rms_epsilon", to_config_value(rms_match.epsilon)); 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 0d7542c7a493..3b668b05bb28 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 @@ -175,6 +175,14 @@ static bool supported_token_count(int64_t token_count) { return token_count >= 1 && token_count <= 2048; } +static int64_t token_capacity_class(int64_t token_count) { + int64_t capacity = 1; + while (capacity < token_count) { + capacity *= 2; + } + return capacity; +} + static bool supported_cache_row_count(int64_t row_count) { return row_count >= 1 && row_count <= 1048576; } @@ -444,7 +452,10 @@ static SetRowsMatch match_set_rows_2d(const Graph & graph, const GraphNode * nod return match; } -static void add_rope_compile_parameters(Dispatch & dispatch, const RopeMatch & match, const char * prefix) { +static void add_rope_compile_parameters(Dispatch & dispatch, + const RopeMatch & match, + const char * prefix, + bool include_input_span = true) { dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_size", to_config_value(match.head_size)); dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".n_dims", to_config_value(match.n_dims)); dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".head_count", to_config_value(match.head_count)); @@ -454,7 +465,10 @@ static void add_rope_compile_parameters(Dispatch & dispatch, const RopeMatch & m to_config_value(match.input_stride1)); dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_stride2", to_config_value(match.input_stride2)); - dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_span", to_config_value(match.input_span)); + if (include_input_span) { + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_span", + to_config_value(match.input_span)); + } dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".mscale", rope_mscale_config_value(match.rope_mscale)); dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".mode", to_config_value(match.mode)); @@ -507,8 +521,12 @@ static Dispatch make_rope_dispatch(const DispatchMatchContext & context, const auto [theta, freq_factors] = add_rope_frequency_bindings(context, match, dispatch_match, "common.rope"); Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kRopeF32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - add_rope_compile_parameters(dispatch, match, "ggml.rope_f32"); + dispatch.kernel.integer_parameters.emplace("input_span", match.input_span); + add_rope_compile_parameters(dispatch, match, "ggml.rope_f32", false); + dispatch.kernel.compile_parameters["ggml.rope_f32.token_capacity"] = + to_config_value(token_capacity_class(match.token_count)); dispatch.bindings.push_back({ match.positions->id, 0, match.positions->byte_count }); dispatch.bindings.push_back({ match.input->storage_root, match.input->storage_offset, match.input_span_bytes }); dispatch.bindings.push_back({ theta, 0, match.theta_bytes }); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp index a5dff78f70aa..e7cd7f67f91c 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp @@ -220,6 +220,7 @@ static bool match_scale_f32_dispatch(const DispatchMatchContext & context, Dispa Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kScaleF32Kernel); + dispatch.kernel.workload_specialization = WorkloadSpecialization::Dynamic; dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); dispatch.kernel.compile_parameters.emplace("ggml.scale_f32.scale", format_float_config(params->scale)); dispatch.kernel.compile_parameters.emplace("ggml.scale_f32.bias", format_float_config(params->bias)); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-symmetric-i4.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-symmetric-i4.h new file mode 100644 index 000000000000..f619d9c07d8b --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-symmetric-i4.h @@ -0,0 +1,37 @@ +#pragma once + +#include "../dispatch-registry.h" +#include "ggml.h" + +namespace ggml::hrx { + +inline bool common_symmetric_i4_lowrow_route_enabled() { + // Disabled because the route causes a substantial numerical performance regression. + return false; +} + +inline bool common_symmetric_i4_lowrow_mul_mat_eligible(const Graph & graph, + const GraphNode & consumer, + const Value & input) { + if (!common_symmetric_i4_lowrow_route_enabled() || !graph.has_index() || consumer.op != GGML_OP_MUL_MAT || + consumer.inputs.size() != 2 || consumer.inputs[1] != input.id || input.type != GGML_TYPE_F32 || + !input.contiguous || input.ne[0] < 256 || input.ne[0] > 32768 || input.ne[0] % 64 != 0 || + input.element_count <= 0 || input.element_count % input.ne[0] != 0) { + return false; + } + + const int64_t token_count = input.element_count / input.ne[0]; + if (token_count < 1 || token_count > 16) { + return false; + } + + const Value * weight = graph.values().find(consumer.inputs[0]); + const Value * output = graph.values().find(consumer.output); + return weight != nullptr && output != nullptr && + (weight->type == GGML_TYPE_Q5_K || weight->type == GGML_TYPE_IQ4_XS) && + weight->alias_source.value < 0 && weight->contiguous && output->contiguous && + output->type == GGML_TYPE_F32 && weight->ne[0] == input.ne[0] && weight->ne[1] == output->ne[0] && + output->ne[0] % 64 == 0 && output->ne[1] == token_count && output->ne[2] == 1 && output->ne[3] == 1; +} + +} // namespace ggml::hrx 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 4622b109141e..03cbc2b8b232 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp @@ -106,6 +106,7 @@ static bool match_unary_f32_dispatch(const DispatchMatchContext & context, Dispa Dispatch dispatch; dispatch.kernel = make_kernel_specialization(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))); diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.cpp new file mode 100644 index 000000000000..1791bdac01fe --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.cpp @@ -0,0 +1,91 @@ +#define GGML_COMMON_IMPL_CPP +#include "iq-codebook-resource.h" + +#include "ggml-common.h" + +#include +#include + +namespace ggml::hrx { + +namespace { + +const std::vector & iq1_grid_bytes() { + static_assert(sizeof(iq1s_grid_gpu) == 8192); + static const std::vector data = [] { + std::vector result; + result.reserve(16384); + for (uint32_t packed : iq1s_grid_gpu) { + for (int lane = 0; lane < 8; ++lane) { + const int shift = lane < 4 ? 8 * lane : 8 * (lane - 4) + 4; + const int8_t value = static_cast(((packed >> shift) & 0xf) - 1); + result.push_back(static_cast(value)); + } + } + return result; + }(); + return data; +} + +const std::vector & iq2_xxs_grid_bytes() { + static_assert(sizeof(iq2xxs_grid) == 2048); + static const std::vector data(reinterpret_cast(iq2xxs_grid), + reinterpret_cast(iq2xxs_grid) + sizeof(iq2xxs_grid)); + return data; +} + +const std::vector & iq2_xs_grid_bytes() { + static_assert(sizeof(iq2xs_grid) == 4096); + static const std::vector data(reinterpret_cast(iq2xs_grid), + reinterpret_cast(iq2xs_grid) + sizeof(iq2xs_grid)); + return data; +} + +const std::vector & iq3_xxs_grid_bytes() { + static_assert(sizeof(iq3xxs_grid) == 1024); + static const std::vector data(reinterpret_cast(iq3xxs_grid), + reinterpret_cast(iq3xxs_grid) + sizeof(iq3xxs_grid)); + return data; +} + +const std::vector & iq2_s_grid_bytes() { + static_assert(sizeof(iq2s_grid) == 8192); + static const std::vector data(reinterpret_cast(iq2s_grid), + reinterpret_cast(iq2s_grid) + sizeof(iq2s_grid)); + return data; +} + +const std::vector & iq3_s_grid_bytes() { + static_assert(sizeof(iq3s_grid) == 2048); + static const std::vector data(reinterpret_cast(iq3s_grid), + reinterpret_cast(iq3s_grid) + sizeof(iq3s_grid)); + return data; +} + +} // namespace + +const IQCodebookResourceSpec & iq_codebook_resource_spec(IQCodebookResource resource) { + static const IQCodebookResourceSpec iq1 = { "common.iq1.grid.v1", iq1_grid_bytes(), 256 }; + static const IQCodebookResourceSpec iq2_xxs = { "common.iq2_xxs.grid.v1", iq2_xxs_grid_bytes(), 256 }; + static const IQCodebookResourceSpec iq2_xs = { "common.iq2_xs.grid.v1", iq2_xs_grid_bytes(), 256 }; + static const IQCodebookResourceSpec iq3_xxs = { "common.iq3_xxs.grid.v1", iq3_xxs_grid_bytes(), 256 }; + static const IQCodebookResourceSpec iq2_s = { "common.iq2_s.grid.v1", iq2_s_grid_bytes(), 256 }; + static const IQCodebookResourceSpec iq3_s = { "common.iq3_s.grid.v1", iq3_s_grid_bytes(), 256 }; + switch (resource) { + case IQCodebookResource::IQ1: + return iq1; + case IQCodebookResource::IQ2XXS: + return iq2_xxs; + case IQCodebookResource::IQ2XS: + return iq2_xs; + case IQCodebookResource::IQ3XXS: + return iq3_xxs; + case IQCodebookResource::IQ2S: + return iq2_s; + case IQCodebookResource::IQ3S: + return iq3_s; + } + std::abort(); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.h b/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.h new file mode 100644 index 000000000000..c22c1a19ebe7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/iq-codebook-resource.h @@ -0,0 +1,26 @@ +#pragma once + +#include +#include +#include + +namespace ggml::hrx { + +enum class IQCodebookResource { + IQ1, + IQ2XXS, + IQ2XS, + IQ3XXS, + IQ2S, + IQ3S, +}; + +struct IQCodebookResourceSpec { + const char * key; + const std::vector & data; + size_t alignment; +}; + +const IQCodebookResourceSpec & iq_codebook_resource_spec(IQCodebookResource resource); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp index c483308f61dd..193f3dac977e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -17,6 +17,7 @@ namespace { static constexpr size_t kOpCount = static_cast(GGML_OP_COUNT); static const std::vector kEmptyRegistrations; +static const std::vector kEmptyActivationConsumers; static bool valid_root_op(ggml_op op) { return op >= 0 && static_cast(op) < kOpCount; @@ -109,6 +110,14 @@ const std::vector & DispatchRegistry::registrations_for_ro return registrations_by_root_[static_cast(root_op)].ordered; } +const std::vector & +DispatchRegistry::activation_consumers_for_root(ggml_op root_op) const { + if (!valid_root_op(root_op) || activation_consumers_by_root_.empty()) { + return kEmptyActivationConsumers; + } + return activation_consumers_by_root_[static_cast(root_op)]; +} + void DispatchRegistryBuilder::add(DispatchRegistration registration) { if (!valid_root_op(registration.root_op) || registration.matcher == nullptr) { return; @@ -126,10 +135,25 @@ void DispatchRegistryBuilder::add(DispatchRegistration registration) { } } +void DispatchRegistryBuilder::add_activation_consumer(DispatchActivationConsumerRegistration registration) { + if (!valid_root_op(registration.root_op) || registration.accepts == nullptr || + registration.accepted_formats == DispatchActivationInputNone) { + return; + } + if (registry_.activation_consumers_by_root_.empty()) { + registry_.activation_consumers_by_root_.resize(kOpCount); + } + registry_.activation_consumers_by_root_[static_cast(registration.root_op)].push_back( + std::move(registration)); +} + DispatchRegistry DispatchRegistryBuilder::build() { if (registry_.registrations_by_root_.empty()) { registry_.registrations_by_root_.resize(kOpCount); } + if (registry_.activation_consumers_by_root_.empty()) { + registry_.activation_consumers_by_root_.resize(kOpCount); + } for (DispatchRegistry::RegistrationGroup & group : registry_.registrations_by_root_) { sort_registrations(group.fused); sort_registrations(group.single_op); diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h index bf92f9f82e1e..f491890c1bd3 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h @@ -6,11 +6,14 @@ #include "status.h" #include +#include #include #include namespace ggml::hrx { +class DispatchRegistry; + struct DispatchTarget { std::string architecture; }; @@ -26,6 +29,15 @@ enum class DispatchSource { Qwen, }; +enum DispatchActivationConsumerSource : uint32_t { + DispatchActivationConsumerSourceCommon = 1u << 0, + DispatchActivationConsumerSourceLlm = 1u << 1, + DispatchActivationConsumerSourceQwen = 1u << 2, + DispatchActivationConsumerSourceAll = DispatchActivationConsumerSourceCommon | + DispatchActivationConsumerSourceLlm | + DispatchActivationConsumerSourceQwen, +}; + struct DispatchMatchContext { const Graph & graph; const GraphNode * root_node = nullptr; @@ -33,6 +45,48 @@ struct DispatchMatchContext { const std::vector & covered_nodes; const CommandPlan & plan; ValueId next_plan_value; + const DispatchRegistry * registry = nullptr; +}; + +enum DispatchActivationInputFormat : uint32_t { + DispatchActivationInputNone = 0, + DispatchActivationInputF16Row = 1u << 0, + DispatchActivationInputF16K16Major = 1u << 1, + DispatchActivationInputQ8_1X4 = 1u << 2, + DispatchActivationInputSymmetricI4K32 = 1u << 3, +}; + +enum DispatchActivationProducerRequirement : uint32_t { + DispatchActivationProducerRequirementNone = 0, + DispatchActivationProducerRequirementPackedK16Economy = 1u << 0, + DispatchActivationProducerRequirementFactoredTiledF16 = 1u << 1, +}; + +enum DispatchActivationConsumerTrait : uint32_t { + DispatchActivationConsumerTraitNone = 0, + DispatchActivationConsumerTraitQ8PublicationPreferred = 1u << 0, +}; + +enum class DispatchActivationConsumerUse { + BandwidthLimited, + Tiled, + Row, + LowRow, +}; + +using DispatchActivationConsumerPredicate = bool (*)(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input); + +struct DispatchActivationConsumerRegistration { + const char * name = ""; + ggml_op root_op = GGML_OP_NONE; + uint32_t accepted_formats = DispatchActivationInputNone; + DispatchActivationConsumerUse use = DispatchActivationConsumerUse::Row; + uint32_t producer_requirements = DispatchActivationProducerRequirementNone; + DispatchSource source = DispatchSource::Common; + DispatchActivationConsumerPredicate accepts = nullptr; + uint32_t traits = DispatchActivationConsumerTraitNone; }; struct DispatchValueAliasRequest { @@ -90,6 +144,9 @@ class DispatchRegistry { const std::vector & single_op_registrations() const { return single_op_registrations_; } + const std::vector & + activation_consumers_for_root(ggml_op root_op) const; + private: friend class DispatchRegistryBuilder; @@ -101,11 +158,13 @@ class DispatchRegistry { std::vector registrations_by_root_; std::vector single_op_registrations_; + std::vector> activation_consumers_by_root_; }; class DispatchRegistryBuilder { public: void add(DispatchRegistration registration); + void add_activation_consumer(DispatchActivationConsumerRegistration registration); DispatchRegistry build(); private: diff --git a/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.cpp b/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.cpp new file mode 100644 index 000000000000..81d7382f2fe7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.cpp @@ -0,0 +1,138 @@ +#include "immutable-program-resource.h" + +#include +#include +#include + +namespace ggml::hrx { +namespace { + +struct ResourceCandidate { + const CommandPlanTransient * transient = nullptr; + size_t matching_count = 0; +}; + +ResourceCandidate find_resource_candidate(const CommandPlan & plan, + const DispatchMatch & match, + const std::string & key) { + ResourceCandidate result; + auto consider = [&](const CommandPlanTransient & transient) { + if (transient.name == key) { + result.transient = &transient; + ++result.matching_count; + } + }; + for (const CommandPlanTransient & transient : plan.transients) { + consider(transient); + } + for (const CommandPlanTransient & transient : match.transients) { + consider(transient); + } + return result; +} + +const CommandPlanConstantInitialization * find_initialization(const CommandPlan & plan, + const DispatchMatch & match, + ValueId value, + size_t & matching_count) { + const CommandPlanConstantInitialization * result = nullptr; + auto consider = [&](const CommandPlanConstantInitialization & initialization) { + if (initialization.value == value) { + result = &initialization; + ++matching_count; + } + }; + for (const CommandPlanConstantInitialization & initialization : plan.constant_initializations) { + consider(initialization); + } + for (const CommandPlanConstantInitialization & initialization : match.constant_initializations) { + consider(initialization); + } + return result; +} + +size_t count_initializations_by_name(const CommandPlan & plan, const DispatchMatch & match, const std::string & key) { + const auto named = [&](const CommandPlanConstantInitialization & initialization) { + return initialization.name == key; + }; + return static_cast( + std::count_if(plan.constant_initializations.begin(), plan.constant_initializations.end(), named)) + + static_cast( + std::count_if(match.constant_initializations.begin(), match.constant_initializations.end(), named)); +} + +bool valid_alignment(size_t alignment) { + return alignment != 0 && (alignment & (alignment - 1)) == 0; +} + +} // namespace + +ValueId require_immutable_program_resource(const DispatchMatchContext & context, + DispatchMatch & match, + const ImmutableProgramResourceSpec & spec) { + if (spec.key.empty()) { + match.status.log("immutable program resource key must not be empty"); + return {}; + } + if (spec.data.empty()) { + match.status.log("immutable program resource '%s' data must not be empty", spec.key.c_str()); + return {}; + } + if (!valid_alignment(spec.alignment)) { + match.status.log("immutable program resource '%s' alignment %zu must be a non-zero power of two", + spec.key.c_str(), spec.alignment); + return {}; + } + + const ResourceCandidate candidate = find_resource_candidate(context.plan, match, spec.key); + if (candidate.matching_count > 1) { + match.status.log("immutable program resource key '%s' is not unique", spec.key.c_str()); + return {}; + } + if (candidate.transient != nullptr) { + size_t matching_initializations = 0; + const CommandPlanConstantInitialization * initialization = + find_initialization(context.plan, match, candidate.transient->value, matching_initializations); + if (matching_initializations != 1 || count_initializations_by_name(context.plan, match, spec.key) != 1 || + initialization == nullptr || initialization->name != spec.key || initialization->offset != 0 || + candidate.transient->size != spec.data.size() || candidate.transient->alignment != spec.alignment || + initialization->data != spec.data) { + match.status.log("immutable program resource key '%s' conflicts with an existing resource", + spec.key.c_str()); + return {}; + } + return candidate.transient->value; + } + + const bool orphaned_initialization = + std::any_of(context.plan.constant_initializations.begin(), context.plan.constant_initializations.end(), + [&](const CommandPlanConstantInitialization & initialization) { + return initialization.name == spec.key; + }) || + std::any_of( + match.constant_initializations.begin(), match.constant_initializations.end(), + [&](const CommandPlanConstantInitialization & initialization) { return initialization.name == spec.key; }); + if (orphaned_initialization) { + match.status.log("immutable program resource key '%s' has an initialization without a transient", + spec.key.c_str()); + return {}; + } + + if (match.transients.size() > std::numeric_limits::max() - match.completion_counter_requests.size()) { + match.status.log("immutable program resource '%s' cannot allocate a plan value", spec.key.c_str()); + return {}; + } + const size_t value_offset = match.transients.size() + match.completion_counter_requests.size(); + if (context.next_plan_value.value < 0 || value_offset > static_cast(std::numeric_limits::max()) || + context.next_plan_value.value > std::numeric_limits::max() - static_cast(value_offset)) { + match.status.log("immutable program resource '%s' cannot allocate a plan value", spec.key.c_str()); + return {}; + } + + const ValueId value(context.next_plan_value.value + static_cast(value_offset)); + match.transients.push_back({ value, spec.key, spec.data.size(), spec.alignment }); + match.constant_initializations.push_back({ value, spec.key, 0, spec.data }); + return value; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.h b/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.h new file mode 100644 index 000000000000..bbe02863f4b7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/immutable-program-resource.h @@ -0,0 +1,22 @@ +#pragma once + +#include "dispatch-registry.h" + +#include +#include +#include +#include + +namespace ggml::hrx { + +struct ImmutableProgramResourceSpec { + std::string key; + const std::vector & data; + size_t alignment = 256; +}; + +ValueId require_immutable_program_resource(const DispatchMatchContext & context, + DispatchMatch & match, + const ImmutableProgramResourceSpec & spec); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp index a974bc1bee0b..75256bc4ec7d 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp @@ -278,9 +278,9 @@ static bool distinct_storage(const Graph & graph, const Value & lhs, const Value return !graph.values().same_storage(lhs.id, rhs.id); } -static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & context, - const AttentionMatMulMatch & root, - const AttentionSetRowsMatch & set_rows) { +static bool is_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & context, + const AttentionMatMulMatch & root, + const AttentionSetRowsMatch & set_rows) { if (prefill_v_cache_fusion_disabled() || root.weight->type != GGML_TYPE_Q6_K || root.input_size != 1536 || root.output_size != 256 || root.token_count != 64 || set_rows.output_format != 16 || set_rows.cache_row_count != 512 || root.output->kind != ValueKind::Transient || @@ -293,6 +293,16 @@ static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & co return false; } + return true; +} + +static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & context, + const AttentionMatMulMatch & root, + const AttentionSetRowsMatch & set_rows) { + if (!is_q6_f16_prefill_v_cache_cell(context, root, set_rows)) { + return false; + } + const size_t activation_bytes = static_cast(root.input_size * root.token_count) * sizeof(ggml_fp16_t); return find_alternate_value(context.graph, context.plan, root.input->id, GGML_TYPE_F16, activation_bytes) != nullptr; @@ -462,6 +472,26 @@ static AttentionSetRowsMatch match_attention_set_rows(const Graph & graph, return match; } +static bool accepts_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + AttentionMatMulMatch root = match_attention_matmul_any_format(context.graph, &consumer); + if (!root.matched() || root.input->id != input.id || !context.graph.has_index()) { + return false; + } + + std::vector layouts; + const Value * after_projection = nullptr; + const GraphNode * set_rows_node = + find_only_consumer_after_layout_aliases(context, root.output, layouts, after_projection, 2); + if (set_rows_node == nullptr || set_rows_node->op != GGML_OP_SET_ROWS) { + return false; + } + const AttentionSetRowsMatch set_rows = match_attention_set_rows(context.graph, set_rows_node, after_projection); + return set_rows.matched() && set_rows.token_count == root.token_count && + set_rows.output_size == root.output_size && is_q6_f16_prefill_v_cache_cell(context, root, set_rows); +} + static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchContext & context) { AttentionQkvMatch match; AttentionMatMulMatch root = match_attention_matmul_any_format(context.graph, context.root_node); @@ -717,6 +747,15 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D } // namespace void register_llm_attention_qkv_dispatches(DispatchRegistryBuilder & registry) { + registry.add_activation_consumer({ + "llm.attention_v_cache.input.f16_row.prefill", + GGML_OP_MUL_MAT, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Tiled, + DispatchActivationProducerRequirementNone, + DispatchSource::Llm, + accepts_q6_f16_prefill_v_cache_cell, + }); registry.add({ "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp index d82613c864e1..56b3888dc82f 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp @@ -27,7 +27,7 @@ static constexpr KernelCatalogRef kGatedDeltaNetPrefillKernel = static constexpr KernelCatalogRef kGatedDeltaNetRmsNormGateKernel = GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate"); static constexpr KernelCatalogRef kRmsNormGateKernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_f16"); + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_publish"); static constexpr KernelCatalogRef kGatedDeltaNetPrefillProjectionEpilogueKernel = GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_projection_epilogue"); static constexpr KernelCatalogRef kGatedDeltaNetInplaceKernel = @@ -42,8 +42,6 @@ static constexpr KernelCatalogRef kGatedDeltaNetSelectedSnapshotKernel = GGML_HR "loom_libs", "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_epilogue"); static constexpr KernelCatalogRef kGatedDeltaNetSelectedRmsQ8Kernel = GGML_HRX_KERNEL_REF( "loom_libs", "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8"); -static constexpr KernelCatalogRef kRmsNormGateQ8Kernel = - GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_gate_f32_q8_1_x4"); static constexpr KernelCatalogRef kCopyF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f32"); static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma"); @@ -739,6 +737,7 @@ static bool match_gated_delta_net_rmsnorm_gate(const DispatchMatchContext & cont const Dispatch & norm = match.dispatches.front(); if (norm.kernel.kernel_id != kRmsNormGateKernel.id || norm.bindings.size() != 5 || norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.hidden_size") != std::to_string(gdn.width) || + norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.publish_q8") != "0" || norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.f16_output_row_width") != std::to_string(gdn.width * gdn.head_count)) { return false; @@ -789,9 +788,10 @@ static bool match_gated_delta_net_dispatch(const DispatchMatchContext & context, if (match.has_projection_epilogue() && !can_fuse_projection_epilogue) { Dispatch epilogue; epilogue.kernel = make_kernel_specialization(kGatedDeltaNetProjectionEpilogueKernel); + epilogue.kernel.workload_specialization = WorkloadSpecialization::Dynamic; + epilogue.kernel.integer_parameters.emplace("element_count", + match.head_count * match.token_count * match.sequence_count); set_compile_parameter(epilogue.kernel, "llm.gated_delta_net.epilogue_head_count", match.head_count); - set_compile_parameter(epilogue.kernel, "llm.gated_delta_net.epilogue_element_count", - match.head_count * match.token_count * match.sequence_count); set_compile_parameter(epilogue.kernel, "llm.gated_delta_net.epilogue_workgroup_size", 256); epilogue.bindings.push_back({ match.alpha_raw->id, 0, match.alpha_raw->byte_count }); epilogue.bindings.push_back({ match.beta_raw->id, 0, match.beta_raw->byte_count }); @@ -993,11 +993,12 @@ static bool match_gated_delta_net_rmsnorm_q8(const DispatchMatchContext & contex return false; } const Dispatch & norm = match.dispatches.front(); - if (norm.kernel.kernel_id != kRmsNormGateQ8Kernel.id || norm.bindings.size() != 5 || + if (norm.kernel.kernel_id != kRmsNormGateKernel.id || norm.bindings.size() != 5 || norm.bindings[0].value != input->id || norm.kernel.integer_parameters.at("token_count") != gdn.head_count * gdn.token_count || norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.hidden_size") != "128" || norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.gate_op") != "15" || + norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.publish_q8") != "1" || norm.kernel.compile_parameters.at("ggml.rmsnorm_gate_f32.f16_output_row_width") != "0") { return false; } diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp index f731ba5da23e..a7e5a4ddb3eb 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp @@ -565,8 +565,9 @@ static bool try_match_binary_fusion(const Graph & graph, SsmConvCoreMatch & matc static void set_generic_ssm_conv_parameters(KernelSpecialization & kernel, const SsmConvCoreMatch & match) { set_compile_parameter(kernel, "llm.ssm_conv.generic.d_conv", match.d_conv); set_compile_parameter(kernel, "llm.ssm_conv.generic.d_inner", match.d_inner); - set_compile_parameter(kernel, "llm.ssm_conv.generic.n_t", match.token_count); - set_compile_parameter(kernel, "llm.ssm_conv.generic.n_s", match.sequence_count); + kernel.integer_parameters.emplace("n_t", match.token_count); + kernel.integer_parameters.emplace("n_s", match.sequence_count); + kernel.workload_specialization = WorkloadSpecialization::Dynamic; set_compile_parameter(kernel, "llm.ssm_conv.generic.unary_op", unary_kind_config_value(match.unary_op)); set_compile_parameter(kernel, "llm.ssm_conv.generic.workgroup_size", 256); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp index 4453f8f18086..fcd7597bb721 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp @@ -1,6 +1,8 @@ #include "dispatch-qwen-matmul.h" #include "dispatch-llm-shapes.h" +#include "dispatch_registration/common/dispatch-activation-publication.h" +#include "dispatch_registration/common/dispatch-mul-mat-common.h" #include "ggml.h" #include "graph/graph-matcher.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" @@ -318,6 +320,22 @@ static QwenAttentionOutputNextQ8Match match_qwen_attention_output_next_q8(const } // namespace +bool qwen_accepts_q6k_q8_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + if (consumer.op != GGML_OP_MUL_MAT || consumer.inputs.size() != 2 || consumer.inputs[1] != input.id) { + return false; + } + const Value * weight = graph_value(context.graph, consumer.inputs[0]); + const Value * output = graph_value(context.graph, consumer.output); + return weight != nullptr && output != nullptr && weight->type == GGML_TYPE_Q6_K && + input.type == GGML_TYPE_F32 && output->type == GGML_TYPE_F32 && weight->contiguous && input.contiguous && + output->contiguous && input.ne[0] == kQwenHiddenSize && input.ne[1] == 1 && input.ne[2] == 1 && + input.ne[3] == 1 && weight->ne[0] == kQwenHiddenSize && weight->ne[1] == kQwenVocabularyCount && + weight->ne[2] == 1 && weight->ne[3] == 1 && output->ne[0] == kQwenVocabularyCount && output->ne[1] == 1 && + output->ne[2] == 1 && output->ne[3] == 1; +} + static void build_qwen_matmul_dispatch(const QwenMatmulMatch & match, DispatchMatch & dispatch_match, size_t root_index) { @@ -354,9 +372,31 @@ static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchCont return false; } + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.normalized_output, capabilities); + const CommonActivationPublicationDemand q8_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (!q8_demand.matched()) { + return false; + } + const ValueId completion_counter(context.next_plan_value.value); - const ValueId q8_output(context.next_plan_value.value + 1); - const size_t q8_output_bytes = q8_1_x4_byte_count(match.token_count, match.output_size); + dispatch_match.completion_counter_requests.push_back({ + completion_counter, + "qwen.decode.attention_output.completion_counter", + 1, + }); + CommonActivationPublication q8_publication; + constexpr const char * q8_publication_name = "qwen.decode.attention_output.next_q8_output"; + if (!common_reserve_activation_publication(context, dispatch_match, *match.normalized_output, q8_demand, + q8_publication_name, q8_publication) || + !common_append_activation_publication(dispatch_match, q8_publication, q8_publication_name)) { + return false; + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kQwenDenseLinearQ4KQ8NextQ8Kernel); @@ -375,25 +415,9 @@ static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchCont dispatch.bindings.push_back({ match.norm_weight->id, 0, match.norm_weight->byte_count }); dispatch.bindings.push_back({ match.normalized_output->id, 0, match.normalized_output->byte_count }); dispatch.bindings.push_back({ completion_counter, 0, sizeof(int32_t) }); - dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + dispatch.bindings.push_back(q8_publication.binding()); dispatch_match.value_aliases.push_back({ match.residual_input->id, match.residual_output->id }); - dispatch_match.completion_counter_requests.push_back({ - completion_counter, - "qwen.decode.attention_output.completion_counter", - 1, - }); - dispatch_match.transients.push_back( - { q8_output, "qwen.decode.attention_output.next_q8_output", q8_output_bytes, 256 }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { match.normalized_output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, - "qwen.decode.attention_output.next_q8_output" }, - metadata_status)) { - dispatch_match.status.append(metadata_status); - return false; - } - if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, dispatch_match.covered_nodes) || (match.projection_get_rows != nullptr && @@ -415,6 +439,15 @@ static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchCont } void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry) { + registry.add_activation_consumer({ + "qwen.matmul.input.q8_1_x4.q6k_vocab", + GGML_OP_MUL_MAT, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Qwen, + qwen_accepts_q6k_q8_input, + }); registry.add({ "qwen.matmul.attention_output_q4k_q8_1_x4_next_q8", GGML_OP_MUL_MAT, diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.h index 59eb88679ff3..28f348e9cacb 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.h @@ -4,6 +4,10 @@ namespace ggml::hrx { +bool qwen_accepts_q6k_q8_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input); + void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry); } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp index db6b45794dbd..dce10fa60c64 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp @@ -1,6 +1,8 @@ #include "dispatch-qwen-rmsnorm.h" #include "dispatch-llm-profiles.h" +#include "../common/dispatch-activation-publication.h" +#include "../common/dispatch-mul-mat-common.h" #include "ggml.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" @@ -46,18 +48,6 @@ static bool is_supported_token_count(int64_t token_count) { return is_llm_supported_query_length(kQwen30BMoeDispatchProfile, token_count); } -static bool has_decode_q8_consumer(const Graph & graph, ValueId value) { - if (!graph.has_index()) { - return false; - } - for (const GraphNode * consumer : graph.index().consumers(value)) { - if (consumer != nullptr && (consumer->op == GGML_OP_MUL_MAT || consumer->op == GGML_OP_MUL_MAT_ID)) { - return true; - } - } - return false; -} - static size_t q8_1_x4_byte_count(int64_t token_count, int64_t hidden_size) { if (token_count <= 0 || hidden_size <= 0) { return 0; @@ -330,17 +320,26 @@ static bool match_qwen_decode_rmsnorm_f32_quantize_q8_1_x4_dispatch(const Dispat context.covered_nodes[rms_match.mul_node_index]) { return false; } - if (!has_decode_q8_consumer(context.graph, rms_match.output->id)) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *rms_match.output, capabilities); + const CommonActivationPublicationDemand demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (!demand.matched()) { return false; } - const size_t q8_byte_count = q8_1_x4_byte_count(rms_match.token_count, rms_match.hidden_size); - if (q8_byte_count == 0) { + CommonActivationPublication publication; + constexpr const char * publication_name = "qwen.decode.q8_hidden"; + if (!common_reserve_activation_publication(context, match, *rms_match.output, demand, publication_name, + publication) || + !common_append_activation_publication(match, publication, publication_name)) { return false; } - const ValueId q8_value = context.next_plan_value; - Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kQwenRmsNormF32QuantizeQ8_1X4Kernel); dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); @@ -353,20 +352,11 @@ static bool match_qwen_decode_rmsnorm_f32_quantize_q8_1_x4_dispatch(const Dispat dispatch.bindings.push_back({ rms_match.input->id, 0, rms_match.input->byte_count }); dispatch.bindings.push_back({ rms_match.weight->id, 0, rms_match.weight->byte_count }); dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); - dispatch.bindings.push_back({ q8_value, 0, q8_byte_count }); - - Status metadata_status; - if (!match.metadata.append_alternate_value( - { rms_match.output->id, q8_value, GGML_TYPE_Q8_1, q8_byte_count, "qwen.decode.q8_hidden" }, - metadata_status)) { - match.status.append(metadata_status); - return false; - } + dispatch.bindings.push_back(publication.binding()); match.covered_nodes.push_back(rms_match.rms_node_index); match.covered_nodes.push_back(rms_match.mul_node_index); match.dispatches.push_back(std::move(dispatch)); - match.transients.push_back({ q8_value, "qwen.decode.q8_hidden", q8_byte_count, 256 }); return match.status.success(); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp index 0a2a63b046d4..ef1d99e35ec2 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp @@ -1,6 +1,8 @@ #include "dispatch-routed-ffn.h" #include "dispatch-llm-shapes.h" +#include "dispatch_registration/common/dispatch-activation-publication.h" +#include "dispatch_registration/common/dispatch-mul-mat-common.h" #include "dispatch_registration/common/dispatch-mul-mat-weight-format.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -32,17 +34,19 @@ static constexpr KernelCatalogRef kQwenRoutedDownWeightedReduceF16F32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_weighted_reduce_f16_f32"); static constexpr KernelCatalogRef kQwenRoutedDownWeightedReduceNextRmsNormF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"); +static constexpr KernelCatalogRef kQwenRoutedDownWeightedReduceNextRmsNormF32PublishQ8Kernel = + GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8"); static constexpr const LlmMoeDispatchProfile & kRoutedFfnProfile = kActiveLlmMoeDispatchProfile; static constexpr int64_t kRoutedFfnInputSize = kRoutedFfnProfile.hidden_size; static constexpr int64_t kRoutedFfnExpertHiddenSize = kRoutedFfnProfile.expert_hidden_size; static constexpr int64_t kRoutedFfnExpertCount = kRoutedFfnProfile.expert_count; static constexpr int64_t kRoutedFfnRouteCount = kRoutedFfnProfile.route_count; -static constexpr size_t kRoutedFfnPlanTransientAlignment = 256; static constexpr const char * kRoutedFfnF16GateUpOutputName = "qwen.moe.gate_up_swiglu_f16"; static constexpr const char * kRoutedFfnF16RoutedDownOutputName = "qwen.moe.routed_down_f16"; static constexpr const char * kRoutedFfnQ8GateUpOutputName = "qwen.decode.moe.gate_up_swiglu_q8"; static constexpr const char * kRoutedFfnQ8HiddenOutputName = "qwen.decode.moe.hidden_q8"; +static constexpr const char * kRoutedFfnQ8NextRmsNormOutputName = "qwen.moe.next_rmsnorm_q8"; static const Value * graph_value(const Graph & graph, ValueId id) { return graph.values().find(id); @@ -61,6 +65,51 @@ static bool same_shape(const Value & lhs, const Value & rhs) { return true; } +static bool is_routed_ffn_down_weight(const Value & value); +static bool is_routed_ffn_projection_output(const Value & value, int64_t token_count); +static bool is_routed_ffn_down_output(const Value & value, int64_t token_count); + +static bool accepts_routed_down_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input, + bool require_q4) { + if (consumer.op != GGML_OP_MUL_MAT_ID || consumer.inputs.size() != 3 || consumer.inputs[1] != input.id) { + return false; + } + const Value * weight = graph_value(context.graph, consumer.inputs[0]); + const Value * output = graph_value(context.graph, consumer.output); + return weight != nullptr && output != nullptr && (!require_q4 || weight->type == GGML_TYPE_Q4_K) && + (weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_Q6_K) && is_routed_ffn_down_weight(*weight) && + is_routed_ffn_projection_output(input, input.ne[2]) && is_routed_ffn_down_output(*output, input.ne[2]) && + find_single_consumer_with_op_through_layout_aliases(context.graph, output->id, GGML_OP_MUL) != nullptr; +} + +static bool accepts_f16_routed_down_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return accepts_routed_down_input(context, consumer, input, false); +} + +static bool accepts_q8_routed_down_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + return accepts_routed_down_input(context, consumer, input, true); +} + +static bool accepts_f16_weighted_reduce_input(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + if (consumer.op != GGML_OP_MUL || consumer.inputs.size() != 2 || + (consumer.inputs[0] != input.id && consumer.inputs[1] != input.id)) { + return false; + } + const ValueId weights_id = consumer.inputs[0] == input.id ? consumer.inputs[1] : consumer.inputs[0]; + const Value * weights = graph_value(context.graph, weights_id); + const Value * output = graph_value(context.graph, consumer.output); + return weights != nullptr && output != nullptr && weights->type == GGML_TYPE_F32 && weights->contiguous && + is_shape(*weights, 1, kRoutedFfnRouteCount, input.ne[2], 1) && same_shape(*output, input); +} + static bool is_profile_rms_norm_epsilon(float eps) { const float expected = kRoutedFfnProfile.rms_norm_epsilon; return eps >= expected * 0.9f && eps <= expected * 1.1f; @@ -377,6 +426,48 @@ static bool match_same_route_projection(const Graph & graph, is_routed_ffn_projection_output(*output, token_count); } +static bool accepts_q8_routed_gate_up_input_impl(const DispatchMatchContext & context, + const GraphNode & consumer, + const Value & input) { + if (consumer.op != GGML_OP_MUL_MAT_ID || consumer.inputs.size() != 3 || consumer.inputs[1] != input.id || + input.type != GGML_TYPE_F32 || !input.contiguous || !is_shape(input, kRoutedFfnInputSize, 1, 1, 1)) { + return false; + } + + const Value * route_ids = graph_value(context.graph, consumer.inputs[2]); + const Value * weight = nullptr; + const Value * output = nullptr; + if (route_ids == nullptr || route_ids->type != GGML_TYPE_I32 || + !is_shape(*route_ids, kRoutedFfnRouteCount, 1, 1, 1) || + !match_same_route_projection(context.graph, consumer, input.id, route_ids->id, 1, weight, output)) { + return false; + } + + const GraphNode * glu_node = find_consumer_with_op(context.graph, output->id, GGML_OP_GLU); + if (glu_node == nullptr || glu_node->inputs.size() != 2 || !is_swiglu_params(glu_node->params)) { + return false; + } + const ValueId other_output = glu_node->inputs[0] == output->id ? glu_node->inputs[1] : + glu_node->inputs[1] == output->id ? glu_node->inputs[0] : ValueId(); + const GraphNode * other_node = producer_with_op(context.graph, other_output, GGML_OP_MUL_MAT_ID); + const Value * other_weight = nullptr; + const Value * other_projection = nullptr; + if (other_node == nullptr || other_node == &consumer || + !match_same_route_projection(context.graph, *other_node, input.id, route_ids->id, 1, other_weight, + other_projection) || + !same_shape(*output, *other_projection)) { + return false; + } + + const Value * glu_output = graph_value(context.graph, glu_node->output); + const GraphNode * down_node = glu_output != nullptr ? + find_single_consumer_with_op(context.graph, glu_output->id, GGML_OP_MUL_MAT_ID) : + nullptr; + return glu_output != nullptr && glu_output->kind == ValueKind::Transient && + is_routed_ffn_projection_output(*glu_output, 1) && down_node != nullptr && down_node->inputs.size() == 3 && + down_node->inputs[2] == route_ids->id && accepts_routed_down_input(context, *down_node, *glu_output, false); +} + static bool match_same_fast_route_projection(const Graph & graph, const GraphNode & node, ValueId expected_input, @@ -423,7 +514,8 @@ static RoutedDownMatch match_routed_ffn_down_grouped(const DispatchMatchContext } const CommandPlanAlternateValue * input_alternate = - find_alternate_value(context.plan, input->id, GGML_TYPE_F16, f16_gate_up_output_size(token_count)); + find_alternate_value(context.graph, context.plan, input->id, GGML_TYPE_F16, + f16_gate_up_output_size(token_count)); if (input_alternate == nullptr) { return {}; } @@ -902,7 +994,8 @@ static WeightedReduceMatch match_routed_ffn_down_weighted_reduce(const DispatchM } const int64_t token_count = match.token_count; const CommandPlanAlternateValue * routed_alternate = - find_alternate_value(context.plan, routed_output->id, GGML_TYPE_F16, f16_routed_down_output_size(token_count)); + find_alternate_value(context.graph, context.plan, routed_output->id, GGML_TYPE_F16, + f16_routed_down_output_size(token_count)); if (routed_alternate == nullptr) { return {}; } @@ -940,14 +1033,18 @@ static bool match_routed_ffn_gate_up_swiglu_f16_wmma_dispatch(const DispatchMatc const std::string weight_format_config = to_config_value(common_mul_mat_format_config_value(weight_format)); const std::string config_prefix = "ggml.mul_mat_id_swiglu_f16_f16"; - const ValueId f16_output(context.next_plan_value.value); - const size_t f16_output_bytes = f16_gate_up_output_size(match.token_count); - dispatch_match.transients.push_back( - { f16_output, kRoutedFfnF16GateUpOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { match.glu_output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16GateUpOutputName }, - metadata_status)) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputF16Row, + DispatchActivationInputF16Row, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.glu_output, capabilities); + const CommonActivationPublicationDemand f16_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + CommonActivationPublication f16_publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.glu_output, f16_demand, + kRoutedFfnF16GateUpOutputName, f16_publication) || + !common_append_activation_publication(dispatch_match, f16_publication, kRoutedFfnF16GateUpOutputName)) { return false; } @@ -972,7 +1069,7 @@ static bool match_routed_ffn_gate_up_swiglu_f16_wmma_dispatch(const DispatchMatc { match.routing_bundle->partition_table, 0, match.routing_bundle->partition_table_byte_count }); dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); - dispatch.bindings.push_back({ f16_output, 0, f16_output_bytes }); + dispatch.bindings.push_back(f16_publication.binding()); if (!append_covered_node(context, match.gate_node, dispatch_match) || !append_covered_node(context, match.up_node, dispatch_match) || @@ -990,10 +1087,24 @@ static bool match_decode_routed_ffn_gate_up_swiglu_q4k_q8_dispatch(const Dispatc return false; } - const size_t q8_output_bytes = - q8_1_x4_byte_count(match.token_count * kRoutedFfnRouteCount, kRoutedFfnExpertHiddenSize); - const ValueId q8_output = match.publish_q8 ? ValueId(context.next_plan_value.value) : ValueId(); - const ValueId completion_counters = match.publish_q8 ? ValueId(context.next_plan_value.value + 1) : ValueId(); + CommonActivationPublication q8_publication; + ValueId completion_counters; + if (match.publish_q8) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.glu_output, capabilities); + const CommonActivationPublicationDemand q8_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (!common_reserve_activation_publication(context, dispatch_match, *match.glu_output, q8_demand, + kRoutedFfnQ8GateUpOutputName, q8_publication) || + !common_append_activation_publication(dispatch_match, q8_publication, kRoutedFfnQ8GateUpOutputName)) { + return false; + } + completion_counters = ValueId(context.next_plan_value.value + 1); + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(match.publish_q8 ? kCommonVectorRoutedGateUpSwiGLUPublishQ8Kernel : @@ -1015,21 +1126,12 @@ static bool match_decode_routed_ffn_gate_up_swiglu_q4k_q8_dispatch(const Dispatc if (match.publish_q8) { dispatch.bindings.push_back( { completion_counters, 0, gate_up_completion_counter_count(match.token_count) * sizeof(int32_t) }); - dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + dispatch.bindings.push_back(q8_publication.binding()); dispatch_match.completion_counter_requests.push_back({ completion_counters, "qwen.decode.moe.gate_up_completion_counters", gate_up_completion_counter_count(match.token_count), }); - dispatch_match.transients.push_back( - { q8_output, kRoutedFfnQ8GateUpOutputName, q8_output_bytes, kRoutedFfnPlanTransientAlignment }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { match.glu_output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, kRoutedFfnQ8GateUpOutputName }, - metadata_status)) { - dispatch_match.status.append(metadata_status); - return false; - } } if (!append_covered_node(context, match.gate_node, dispatch_match) || @@ -1119,9 +1221,30 @@ static bool match_decode_routed_ffn_down_next_q8_dispatch(const DispatchMatchCon return false; } + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.reduce.next_rmsnorm.output, capabilities); + const CommonActivationPublicationDemand q8_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (!q8_demand.matched()) { + return false; + } + const ValueId completion_counter_value(context.next_plan_value.value); - const ValueId q8_output(context.next_plan_value.value + 1); - const size_t q8_output_bytes = q8_1_x4_byte_count(match.token_count, kRoutedFfnInputSize); + dispatch_match.completion_counter_requests.push_back({ + completion_counter_value, + "qwen.decode.moe.routed_down_completion_counter", + 1, + }); + CommonActivationPublication q8_publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.reduce.next_rmsnorm.output, q8_demand, + kRoutedFfnQ8HiddenOutputName, q8_publication) || + !common_append_activation_publication(dispatch_match, q8_publication, kRoutedFfnQ8HiddenOutputName)) { + return false; + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(match.kernel); @@ -1150,25 +1273,9 @@ static bool match_decode_routed_ffn_down_next_q8_dispatch(const DispatchMatchCon dispatch.bindings.push_back( { match.reduce.next_rmsnorm.norm_weight->id, 0, match.reduce.next_rmsnorm.norm_weight->byte_count }); dispatch.bindings.push_back({ completion_counter_value, 0, sizeof(int32_t) }); - dispatch.bindings.push_back({ q8_output, 0, q8_output_bytes }); + dispatch.bindings.push_back(q8_publication.binding()); dispatch_match.value_aliases.push_back({ match.reduce.residual_input->id, match.output->id }); - dispatch_match.completion_counter_requests.push_back({ - completion_counter_value, - "qwen.decode.moe.routed_down_completion_counter", - 1, - }); - dispatch_match.transients.push_back( - { q8_output, kRoutedFfnQ8HiddenOutputName, q8_output_bytes, kRoutedFfnPlanTransientAlignment }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { match.reduce.next_rmsnorm.output->id, q8_output, GGML_TYPE_Q8_1, q8_output_bytes, - kRoutedFfnQ8HiddenOutputName }, - metadata_status)) { - dispatch_match.status.append(metadata_status); - return false; - } - if (!append_covered_node(context, context.root_node, dispatch_match) || !append_covered_node(context, match.reduce.weighted_node, dispatch_match)) { return false; @@ -1201,14 +1308,19 @@ static bool build_routed_ffn_down_grouped_dispatch(const DispatchMatchContext & return false; } - const ValueId f16_output(context.next_plan_value.value); - const size_t f16_output_bytes = f16_routed_down_output_size(match.token_count); - dispatch_match.transients.push_back( - { f16_output, kRoutedFfnF16RoutedDownOutputName, f16_output_bytes, kRoutedFfnPlanTransientAlignment }); - Status metadata_status; - if (!dispatch_match.metadata.append_alternate_value( - { match.output->id, f16_output, GGML_TYPE_F16, f16_output_bytes, kRoutedFfnF16RoutedDownOutputName }, - metadata_status)) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputF16Row, + DispatchActivationInputF16Row, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.output, capabilities); + const CommonActivationPublicationDemand f16_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + CommonActivationPublication f16_publication; + if (!common_reserve_activation_publication(context, dispatch_match, *match.output, f16_demand, + kRoutedFfnF16RoutedDownOutputName, f16_publication) || + !common_append_activation_publication(dispatch_match, f16_publication, + kRoutedFfnF16RoutedDownOutputName)) { return false; } @@ -1220,7 +1332,7 @@ static bool build_routed_ffn_down_grouped_dispatch(const DispatchMatchContext & dispatch.bindings.push_back( { match.routing_bundle->expert_table, 0, match.routing_bundle->expert_table_byte_count }); dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); - dispatch.bindings.push_back({ f16_output, 0, f16_output_bytes }); + dispatch.bindings.push_back(f16_publication.binding()); dispatch_match.covered_nodes.push_back(context.root_index); dispatch_match.dispatches.push_back(std::move(dispatch)); @@ -1236,21 +1348,50 @@ static bool match_routed_ffn_down_weighted_reduce_dispatch(const DispatchMatchCo const bool use_next_rmsnorm = match.next_rmsnorm.matched() && match.output->kind == ValueKind::Transient && residual_input_is_safe_for_in_place(context, match); + CommonActivationPublication q8_publication; + if (use_next_rmsnorm) { + const CommonActivationPublicationCapabilities capabilities = { + DispatchActivationInputQ8_1X4, + DispatchActivationInputQ8_1X4, + }; + const CommonActivationPublicationPlan publication_plan = + common_select_activation_publication_plan(context, *match.next_rmsnorm.output, capabilities); + const CommonActivationPublicationDemand q8_demand = + publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; + if (q8_demand.matched() && + (!common_reserve_activation_publication(context, dispatch_match, *match.next_rmsnorm.output, q8_demand, + kRoutedFfnQ8NextRmsNormOutputName, q8_publication) || + !common_append_activation_publication(dispatch_match, q8_publication, + kRoutedFfnQ8NextRmsNormOutputName))) { + return false; + } + } + Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(use_next_rmsnorm ? kQwenRoutedDownWeightedReduceNextRmsNormF32Kernel : - kQwenRoutedDownWeightedReduceF16F32Kernel); + dispatch.kernel = make_kernel_specialization( + q8_publication.matched() ? kQwenRoutedDownWeightedReduceNextRmsNormF32PublishQ8Kernel : + use_next_rmsnorm ? kQwenRoutedDownWeightedReduceNextRmsNormF32Kernel : + kQwenRoutedDownWeightedReduceF16F32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); add_routed_down_compile_parameters(dispatch, match.token_count); if (use_next_rmsnorm) { dispatch_match.value_aliases.push_back({ match.residual_input->id, match.output->id }); dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(kRoutedFfnInputSize)); dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); + if (q8_publication.matched()) { + dispatch.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + to_config_value(match.token_count * + (kRoutedFfnInputSize / 128))); + } dispatch.bindings.push_back({ match.route_weights->id, 0, match.route_weights->byte_count }); dispatch.bindings.push_back({ match.routed_alternate->alternate_value, 0, match.routed_alternate->byte_count }); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch.bindings.push_back( { match.next_rmsnorm.norm_weight->id, 0, match.next_rmsnorm.norm_weight->byte_count }); dispatch.bindings.push_back({ match.next_rmsnorm.output->id, 0, match.next_rmsnorm.output->byte_count }); + if (q8_publication.matched()) { + dispatch.bindings.push_back(q8_publication.binding()); + } } else { dispatch.bindings.push_back({ match.route_weights->id, 0, match.route_weights->byte_count }); dispatch.bindings.push_back({ match.routed_alternate->alternate_value, 0, match.routed_alternate->byte_count }); @@ -1296,6 +1437,42 @@ static bool match_routed_ffn_down_q6k_f16_wmma_grouped_dispatch(const DispatchMa } // namespace void register_routed_ffn_dispatches(DispatchRegistryBuilder & registry) { + registry.add_activation_consumer({ + "llm.routed_ffn.input.q8_1_x4.gate_up", + GGML_OP_MUL_MAT_ID, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Llm, + accepts_q8_routed_gate_up_input_impl, + }); + registry.add_activation_consumer({ + "llm.routed_ffn.input.q8_1_x4.down", + GGML_OP_MUL_MAT_ID, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, + DispatchSource::Llm, + accepts_q8_routed_down_input, + }); + registry.add_activation_consumer({ + "llm.routed_ffn.input.f16_row.down", + GGML_OP_MUL_MAT_ID, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Row, + DispatchActivationProducerRequirementNone, + DispatchSource::Llm, + accepts_f16_routed_down_input, + }); + registry.add_activation_consumer({ + "llm.routed_ffn.input.f16_row.weighted_reduce", + GGML_OP_MUL, + DispatchActivationInputF16Row, + DispatchActivationConsumerUse::Row, + DispatchActivationProducerRequirementNone, + DispatchSource::Llm, + accepts_f16_weighted_reduce_input, + }); registry.add({ "llm.routed_ffn.decode_gate_up_swiglu_q4k_q8", GGML_OP_MUL_MAT_ID, 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 1f6348c02a23..bf1928273c30 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 @@ -2,6 +2,21 @@ "schema": "ggml-hrx-kernel-corpus-v1", "upstream_revision": "local", "files": [ + { + "path": "ops/matmul/vector_codebook_f32_f32.loom" + }, + { + "path": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + }, + { + "path": "motifs/matmul/vector_codebook_activation.loom" + }, + { + "path": "motifs/matmul/vector_codebook_decode_pair.loom" + }, + { + "path": "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom" + }, { "path": "ops/elementwise/unary_f32.loom" }, @@ -11,6 +26,9 @@ { "path": "ops/elementwise/binary_bc_f32.loom" }, + { + "path": "ops/elementwise/binary_bc_publish_q8_1_x4.loom" + }, { "path": "ops/elementwise/scale_f32.loom" }, @@ -44,6 +62,9 @@ { "path": "ops/elementwise/binary_f32.loom" }, + { + "path": "ops/elementwise/binary_publish_q8_1_x4.loom" + }, { "path": "ops/normalization/rmsnorm_f32.loom" }, @@ -63,10 +84,7 @@ "path": "ops/normalization/rmsnorm_binary_f32_k16.loom" }, { - "path": "ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom" - }, - { - "path": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom" + "path": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom" }, { "path": "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" @@ -78,10 +96,7 @@ "path": "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" }, { - "path": "ops/normalization/rmsnorm_gate_publish_f16.loom" - }, - { - "path": "ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom" + "path": "ops/normalization/rmsnorm_gate_publish.loom" }, { "path": "ops/normalization/rmsnorm_mul_rope_f32.loom" @@ -191,6 +206,21 @@ { "path": "motifs/matmul/q6_q8_1_x4_decode.loom" }, + { + "path": "motifs/matmul/q1_0_q8_1_x4_decode.loom" + }, + { + "path": "ops/matmul/q1_0_q8_1_x4_packed_test.loom" + }, + { + "path": "ops/matmul/vector_q1_0_q8_1_x4_f32.loom" + }, + { + "path": "ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom" + }, + { + "path": "ops/matmul/q1_0_vector_schedule_test.loom" + }, { "path": "motifs/routing/router_projection_storage.loom" }, @@ -605,6 +635,9 @@ { "path": "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom" }, + { + "path": "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom" + }, { "path": "motifs/matmul/tiled_postprocess_f32.loom" }, @@ -620,39 +653,77 @@ ], "exports": [ { - "name": "ggml_copy_transpose_f16", + "name": "ggml_mul_mat_vector_iq2_s_f32_f32", "family": "loom_libs", - "symbol": "ggml_copy_transpose_f16", - "source": "ops/copy/transpose_f16.loom", - "workload_parameters": [], - "launch_parameters": [], + "symbol": "ggml_mul_mat_vector_iq2_s_f32_f32", + "source": "ops/matmul/vector_codebook_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], "bindings": [ "input", - "output" + "weight", + "iq2_s_grid", + "output", + "next_output" ], "binding_access": [ "read", + "read", + "read", + "read_write", "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/transpose_f16.loom" + "ops/matmul/vector_codebook_f32_f32.loom" ], "library_sources": [ - "motifs/copy/transpose_tile.loom" + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/copy/transpose_tile.loom" + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_copy_f16_k16_major", + "name": "ggml_mul_mat_vector_iq3_s_f32_f32", "family": "loom_libs", - "symbol": "ggml_copy_f16_k16_major", - "source": "ops/copy/copy_f16_k16_major.loom", + "symbol": "ggml_mul_mat_vector_iq3_s_f32_f32", + "source": "ops/matmul/vector_codebook_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -667,89 +738,145 @@ ], "bindings": [ "input", - "output" + "weight", + "iq3_s_grid", + "output", + "next_output" ], "binding_access": [ "read", + "read", + "read", + "read_write", "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/copy_f16_k16_major.loom" + "ops/matmul/vector_codebook_f32_f32.loom" ], "library_sources": [ - "motifs/copy/layout_k16_major.loom", - "motifs/copy/convert_f32_f16.loom" + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/copy/layout_k16_major.loom", - "motifs/copy/convert_f32_f16.loom" + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_copy_f32_f16", + "name": "ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", "family": "loom_libs", - "symbol": "ggml_copy_f32_f16", - "source": "ops/copy/copy_f32_f16.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", + "source": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom", "workload_parameters": [ { - "name": "element_count", + "name": "token_count", "type": "index" } ], "launch_parameters": [ { - "name": "element_count", + "name": "token_count", "type": "index" } ], "bindings": [ - "source", + "input", + "weight", + "iq2_s_grid", "output" ], "binding_access": [ "read", - "write" + "read", + "read", + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/copy_f32_f16.loom" + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" ], "library_sources": [ - "motifs/copy/convert_f32_f16.loom" + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/copy/convert_f32_f16.loom" + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_copy_f32", + "name": "ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32", "family": "loom_libs", - "symbol": "ggml_copy_f32", - "source": "ops/copy/copy_f32.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32", + "source": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom", "workload_parameters": [ { - "name": "element_count", + "name": "token_count", "type": "index" } ], "launch_parameters": [ { - "name": "element_count", + "name": "token_count", "type": "index" } ], "bindings": [ - "source", + "input", + "weight", + "iq3_s_grid", "output" ], "binding_access": [ + "read", + "read", "read", "read_write" ], @@ -757,77 +884,107 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/copy_f32.loom" + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" ], "library_sources": [ - "motifs/copy/vector_load_store.loom" + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/copy/vector_load_store.loom" + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_copy_strided_source_f32", + "name": "ggml_copy_transpose_f16", "family": "loom_libs", - "symbol": "ggml_copy_strided_source_f32", - "source": "ops/copy/strided_source_f32.loom", + "symbol": "ggml_copy_transpose_f16", + "source": "ops/copy/transpose_f16.loom", "workload_parameters": [ { - "name": "element_count", + "name": "row_count", "type": "index" }, { - "name": "source_span", + "name": "column_count", "type": "index" } ], "launch_parameters": [ { - "name": "element_count", + "name": "row_count", "type": "index" }, { - "name": "source_span", + "name": "column_count", "type": "index" } ], "bindings": [ - "source", + "input", "output" ], "binding_access": [ "read", - "read_write" + "write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/strided_source_f32.loom" + "ops/copy/transpose_f16.loom" ], "library_sources": [ - "motifs/copy/strided_source.loom" + "motifs/copy/transpose_tile.loom" ] }, "compile_dependencies": [ - "motifs/copy/strided_source.loom" + "motifs/copy/transpose_tile.loom" ] }, { - "name": "ggml_concat_dim0_f32", + "name": "ggml_copy_f16_k16_major", "family": "loom_libs", - "symbol": "ggml_concat_dim0_f32", - "source": "ops/copy/concat_dim0_f32.loom", - "workload_parameters": [], - "launch_parameters": [], + "symbol": "ggml_copy_f16_k16_major", + "source": "ops/copy/copy_f16_k16_major.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], "bindings": [ - "lhs", - "rhs", + "input", "output" ], "binding_access": [ - "read", "read", "write" ], @@ -835,34 +992,218 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/copy/concat_dim0_f32.loom" + "ops/copy/copy_f16_k16_major.loom" ], "library_sources": [ - "motifs/copy/concat_dim0.loom" + "motifs/copy/layout_k16_major.loom", + "motifs/copy/convert_f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/copy/concat_dim0.loom" + "motifs/copy/layout_k16_major.loom", + "motifs/copy/convert_f32_f16.loom" ] }, { - "name": "ggml_concat_dim0_strided_source_f32", + "name": "ggml_copy_f32_f16", "family": "loom_libs", - "symbol": "ggml_concat_dim0_strided_source_f32", - "source": "ops/copy/concat_dim0_strided_source_f32.loom", + "symbol": "ggml_copy_f32_f16", + "source": "ops/copy/copy_f32_f16.loom", "workload_parameters": [ { - "name": "lhs_span", - "type": "index" - }, - { - "name": "rhs_span", + "name": "element_count", "type": "index" } ], "launch_parameters": [ { - "name": "lhs_span", + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "source", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/copy_f32_f16.loom" + ], + "library_sources": [ + "motifs/copy/convert_f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/convert_f32_f16.loom" + ] + }, + { + "name": "ggml_copy_f32", + "family": "loom_libs", + "symbol": "ggml_copy_f32", + "source": "ops/copy/copy_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "source", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/copy_f32.loom" + ], + "library_sources": [ + "motifs/copy/vector_load_store.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/vector_load_store.loom" + ] + }, + { + "name": "ggml_copy_strided_source_f32", + "family": "loom_libs", + "symbol": "ggml_copy_strided_source_f32", + "source": "ops/copy/strided_source_f32.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", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/strided_source_f32.loom" + ], + "library_sources": [ + "motifs/copy/strided_source.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/strided_source.loom" + ] + }, + { + "name": "ggml_concat_dim0_f32", + "family": "loom_libs", + "symbol": "ggml_concat_dim0_f32", + "source": "ops/copy/concat_dim0_f32.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + }, + { + "name": "lhs_span", + "type": "index" + }, + { + "name": "rhs_span", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + }, + { + "name": "lhs_span", + "type": "index" + }, + { + "name": "rhs_span", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/concat_dim0.loom" + ] + }, + { + "name": "ggml_concat_dim0_strided_source_f32", + "family": "loom_libs", + "symbol": "ggml_concat_dim0_strided_source_f32", + "source": "ops/copy/concat_dim0_strided_source_f32.loom", + "workload_parameters": [ + { + "name": "lhs_span", + "type": "index" + }, + { + "name": "rhs_span", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "lhs_span", "type": "index" }, { @@ -1095,6 +1436,51 @@ "motifs/elementwise/binary_apply.loom" ] }, + { + "name": "ggml_binary_f32_publish_q8_1_x4", + "family": "loom_libs", + "symbol": "ggml_binary_f32_publish_q8_1_x4", + "source": "ops/elementwise/binary_publish_q8_1_x4.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/binary_publish_q8_1_x4.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, { "name": "ggml_binary_swiglu_symmetric_i4_k32", "family": "loom_libs", @@ -1224,25 +1610,94 @@ ] }, { - "name": "ggml_rmsnorm_f32", + "name": "ggml_binary_bc_f32_publish_q8_1_x4", "family": "loom_libs", - "symbol": "ggml_rmsnorm_f32", - "source": "ops/normalization/rmsnorm_f32.loom", + "symbol": "ggml_binary_bc_f32_publish_q8_1_x4", + "source": "ops/elementwise/binary_bc_publish_q8_1_x4.loom", "workload_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + }, + { + "name": "src0_element_count", + "type": "index" + }, + { + "name": "src1_element_count", + "type": "index" } ], "launch_parameters": [ { "name": "token_count", "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + }, + { + "name": "src0_element_count", + "type": "index" + }, + { + "name": "src1_element_count", + "type": "index" } ], "bindings": [ - "input", - "output" + "lhs", + "rhs", + "output", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/binary_bc_publish_q8_1_x4.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_rmsnorm_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_f32", + "source": "ops/normalization/rmsnorm_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" ], "binding_access": [ "read", @@ -1529,61 +1984,10 @@ ] }, { - "name": "ggml_rmsnorm_binary_q8_1_x4", - "family": "loom_libs", - "symbol": "ggml_rmsnorm_binary_q8_1_x4", - "source": "ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "rhs", - "output", - "q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom" - ], - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm_binary.loom", - "motifs/normalization/publish_q8.loom", - "motifs/elementwise/binary_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "motifs/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm_binary.loom", - "motifs/normalization/publish_q8.loom", - "motifs/elementwise/binary_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - { - "name": "ggml_rmsnorm_binary_q8_1_x4_f16", + "name": "ggml_rmsnorm_binary_q8_1_x4_publish", "family": "loom_libs", - "symbol": "ggml_rmsnorm_binary_q8_1_x4_f16", - "source": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom", + "symbol": "ggml_rmsnorm_binary_q8_1_x4_publish", + "source": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom", "workload_parameters": [ { "name": "token_count", @@ -1614,7 +2018,7 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom" + "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom" ], "library_sources": [ "motifs/normalization/rmsnorm.loom", @@ -1786,65 +2190,10 @@ ] }, { - "name": "ggml_rmsnorm_gate_f32_f16", - "family": "loom_libs", - "symbol": "ggml_rmsnorm_gate_f32_f16", - "source": "ops/normalization/rmsnorm_gate_publish_f16.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "raw_gate", - "output", - "f16_output" - ], - "binding_access": [ - "read", - "read", - "read", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ops/normalization/rmsnorm_gate_publish_f16.loom" - ], - "library_sources": [ - "motifs/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm_gate.loom", - "motifs/normalization/publish_f16.loom", - "motifs/normalization/publish_q8.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "motifs/normalization/rmsnorm.loom", - "motifs/normalization/rmsnorm_gate.loom", - "motifs/normalization/publish_f16.loom", - "motifs/normalization/publish_q8.loom", - "motifs/elementwise/unary_apply.loom", - "motifs/quantize_q8_1_x4.loom" - ] - }, - { - "name": "ggml_rmsnorm_gate_f32_q8_1_x4", + "name": "ggml_rmsnorm_gate_f32_publish", "family": "loom_libs", - "symbol": "ggml_rmsnorm_gate_f32_q8_1_x4", - "source": "ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom", + "symbol": "ggml_rmsnorm_gate_f32_publish", + "source": "ops/normalization/rmsnorm_gate_publish.loom", "workload_parameters": [ { "name": "token_count", @@ -1862,7 +2211,7 @@ "weight", "raw_gate", "output", - "q8_output" + "alternate_output" ], "binding_access": [ "read", @@ -1875,7 +2224,7 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom" + "ops/normalization/rmsnorm_gate_publish.loom" ], "library_sources": [ "motifs/normalization/rmsnorm.loom", @@ -2157,6 +2506,137 @@ "motifs/elementwise/binary_apply.loom" ] }, + { + "name": "ggml_mul_mat_swiglu_q1_0_q8_1_x4_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_q1_0_q8_1_x4_f32", + "source": "ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom" + ], + "library_sources": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_q1_0_q8_1_x4_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_q1_0_q8_1_x4_f32", + "source": "ops/matmul/vector_q1_0_q8_1_x4_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_q1_0_q8_1_x4_f32.loom" + ], + "library_sources": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_q1_0_q8_1_x4_rows2_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_q1_0_q8_1_x4_rows2_f32", + "source": "ops/matmul/vector_q1_0_q8_1_x4_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_q1_0_q8_1_x4_f32.loom" + ], + "library_sources": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/q1_0_q8_1_x4_decode.loom", + "motifs/quantize_q8_1_x4_group.loom" + ] + }, { "name": "ggml_mul_mat_vector_f32_f32", "family": "loom_libs", @@ -6509,13 +6989,21 @@ { "name": "token_count", "type": "index" + }, + { + "name": "input_span", + "type": "index" } ], "launch_parameters": [ { "name": "token_count", "type": "index" - } + }, + { + "name": "input_span", + "type": "index" + } ], "bindings": [ "positions", @@ -6801,8 +7289,26 @@ "family": "loom_libs", "symbol": "llm_ssm_conv_f32", "source": "ops/ssm_conv/generic_f32.loom", - "workload_parameters": [], - "launch_parameters": [], + "workload_parameters": [ + { + "name": "n_t", + "type": "index" + }, + { + "name": "n_s", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "n_t", + "type": "index" + }, + { + "name": "n_s", + "type": "index" + } + ], "bindings": [ "window", "filter", @@ -6867,8 +7373,26 @@ "family": "loom_libs", "symbol": "llm_ssm_conv_binary_f32", "source": "ops/ssm_conv/generic_binary_f32.loom", - "workload_parameters": [], - "launch_parameters": [], + "workload_parameters": [ + { + "name": "n_t", + "type": "index" + }, + { + "name": "n_s", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "n_t", + "type": "index" + }, + { + "name": "n_s", + "type": "index" + } + ], "bindings": [ "window", "filter", @@ -7692,45 +8216,721 @@ "family": "loom_libs", "symbol": "llm_gated_delta_net_projection_epilogue_f32", "source": "ops/gated_delta_net/projection_epilogue_f32.loom", - "workload_parameters": [], - "launch_parameters": [], + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "gate_dst", + "beta_dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/projection_epilogue_f32.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/publish_f32.loom" + ] + }, + { + "name": "ggml_mul_mat_skinny_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/skinny_postprocess_f32.loom", + "motifs/matmul/skinny_publish_f32.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32_aligned", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_aligned", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32_token64_aligned", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_token64_aligned", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32_aligned", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32_aligned", + "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/normalization/publish_f16.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/normalization/publish_f16.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], "bindings": [ - "alpha_raw", - "beta_raw", - "bias", - "a_scale", - "gate_dst", - "beta_dst" + "input", + "gate_weight", + "up_weight", + "output" ], "binding_access": [ "read", "read", "read", - "read", - "write", - "write" + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/gated_delta_net/projection_epilogue_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" ], "library_sources": [ - "motifs/gated_delta_net/projection_epilogue.loom", - "motifs/gated_delta_net/publish_f32.loom" + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, "compile_dependencies": [ - "motifs/gated_delta_net/projection_epilogue.loom", - "motifs/gated_delta_net/publish_f32.loom" + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" ] }, { - "name": "ggml_mul_mat_skinny_input_f32_publish_f32", + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", "family": "loom_libs", - "symbol": "ggml_mul_mat_skinny_input_f32_publish_f32", - "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7745,10 +8945,12 @@ ], "bindings": [ "input", - "weight", + "gate_weight", + "up_weight", "output" ], "binding_access": [ + "read", "read", "read", "read_write" @@ -7757,12 +8959,13 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/skinny_input_f32_publish_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7775,9 +8978,10 @@ ] }, "compile_dependencies": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7790,10 +8994,10 @@ ] }, { - "name": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - "source": "ops/matmul/skinny_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7808,7 +9012,8 @@ ], "bindings": [ "input", - "weight", + "gate_weight", + "up_weight", "bias", "residual_input", "residual_output" @@ -7818,18 +9023,20 @@ "read", "read", "read", + "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/skinny_input_f32_publish_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7842,9 +9049,10 @@ ] }, "compile_dependencies": [ - "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/skinny_postprocess_f32.loom", - "motifs/matmul/skinny_publish_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7857,10 +9065,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_publish_f32", + "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32_aligned", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32_aligned", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7875,10 +9083,16 @@ ], "bindings": [ "input", - "weight", - "output" + "gate_weight", + "up_weight", + "bias", + "residual_input", + "residual_output" ], "binding_access": [ + "read", + "read", + "read", "read", "read", "read_write" @@ -7887,10 +9101,11 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -7905,7 +9120,8 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -7920,10 +9136,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "name": "ggml_mul_mat_vector_iq3_xxs_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_vector_iq3_xxs_f32_f32", + "source": "ops/matmul/vector_codebook_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -7939,10 +9155,12 @@ "bindings": [ "input", "weight", + "iq3_xxs_grid", "output", - "f16_output" + "next_output" ], "binding_access": [ + "read", "read", "read", "read_write", @@ -7952,12 +9170,12 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" + "ops/matmul/vector_codebook_f32_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7970,9 +9188,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -7985,10 +9203,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "name": "ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - "source": "ops/matmul/tiled_input_f32_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", + "source": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -8004,25 +9222,23 @@ "bindings": [ "input", "weight", - "bias", - "residual_input", - "residual_output" + "iq3_xxs_grid", + "output" ], "binding_access": [ "read", "read", "read", - "read", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_input_f32_publish_f32.loom" + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8037,7 +9253,7 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8052,10 +9268,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", + "name": "ggml_mul_mat_vector_iq1_s_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom", + "symbol": "ggml_mul_mat_vector_iq1_s_f32_f32", + "source": "ops/matmul/vector_codebook_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -8070,10 +9286,10 @@ ], "bindings": [ "input", - "gate_weight", - "up_weight", + "weight", + "iq1_grid", "output", - "f16_output" + "next_output" ], "binding_access": [ "read", @@ -8086,13 +9302,12 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom" + "ops/matmul/vector_codebook_f32_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/normalization/publish_f16.loom", + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -8105,10 +9320,9 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", - "motifs/matmul/tiled_publish.loom", - "motifs/normalization/publish_f16.loom", + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", "motifs/dequant.loom", "motifs/elementwise/unary_apply.loom", "motifs/q4_k_f16.loom", @@ -8121,10 +9335,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + "name": "ggml_mul_mat_vector_iq1_m_f32_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "symbol": "ggml_mul_mat_vector_iq1_m_f32_f32", + "source": "ops/matmul/vector_codebook_f32_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -8139,8 +9353,75 @@ ], "bindings": [ "input", - "gate_weight", - "up_weight", + "weight", + "iq1_grid", + "output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_codebook_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_codebook_activation.loom", + "motifs/matmul/vector_codebook_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32", + "source": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "iq1_grid", "output" ], "binding_access": [ @@ -8153,11 +9434,10 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8172,8 +9452,7 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8188,10 +9467,10 @@ ] }, { - "name": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", + "name": "ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32", "family": "loom_libs", - "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", - "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "symbol": "ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32", + "source": "ops/matmul/tiled_input_f32_codebook_publish_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -8206,15 +9485,11 @@ ], "bindings": [ "input", - "gate_weight", - "up_weight", - "bias", - "residual_input", - "residual_output" + "weight", + "iq1_grid", + "output" ], "binding_access": [ - "read", - "read", "read", "read", "read", @@ -8224,11 +9499,10 @@ "compile_recipe": { "mode": "direct", "primary_sources": [ - "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" ], "library_sources": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", @@ -8243,8 +9517,7 @@ ] }, "compile_dependencies": [ - "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", - "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom", "motifs/matmul/tiled_postprocess_f32.loom", "motifs/matmul/tiled_publish.loom", "motifs/dequant.loom", diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/concat_dim0.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/concat_dim0.loom index 84f25094f667..c9526fb51851 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/concat_dim0.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/concat_dim0.loom @@ -3,7 +3,6 @@ config.decl @ggml.concat_dim0_f32.lhs_width : %value: index where [range(%value, 1, 65536)] config.decl @ggml.concat_dim0_f32.rhs_width : %value: index where [range(%value, 1, 65536)] -config.decl @ggml.concat_dim0_f32.row_count : %value: index where [range(%value, 1, 1048576)] config.decl @ggml.concat_dim0_strided_source_f32.lhs_width : %value: index where [range(%value, 1, 65536)] config.decl @ggml.concat_dim0_strided_source_f32.rhs_width : %value: index where [range(%value, 1, 65536)] @@ -19,18 +18,19 @@ config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride1 : %value: index whe config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride2 : %value: index where [range(%value, 1, 1073741824)] config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride3 : %value: index where [range(%value, 1, 1073741824)] -template.decl @ggml.copy.concat_dim0.f32_body(%lhs: buffer, %rhs: buffer, %output: buffer) +template.decl @ggml.copy.concat_dim0.f32_body(%row_count: index, %lhs_span: index, %rhs_span: index, %lhs: buffer, %rhs: buffer, %output: buffer) template.decl @ggml.copy.concat_dim0.strided_source_f32_body(%lhs_span: index, %rhs_span: index, %lhs: buffer, %rhs: buffer, %output: buffer) -template.def<@ggml.copy.concat_dim0.f32_body> device @ggml_copy_concat_dim0_f32_body(%lhs: buffer, %rhs: buffer, %output: buffer) { +template.def<@ggml.copy.concat_dim0.f32_body> device @ggml_copy_concat_dim0_f32_body(%row_count0: index, %lhs_span0: index, %rhs_span0: index, %lhs: buffer, %rhs: buffer, %output: buffer) { %base = index.constant 0 : offset %workgroup_size = index.constant 256 : index %lhs_width0 = config.get @ggml.concat_dim0_f32.lhs_width : index %rhs_width0 = config.get @ggml.concat_dim0_f32.rhs_width : index - %row_count0 = config.get @ggml.concat_dim0_f32.row_count : index %lhs_width = index.assume %lhs_width0 [range(%lhs_width0, 1, 65536)] : index %rhs_width = index.assume %rhs_width0 [range(%rhs_width0, 1, 65536)] : index %row_count = index.assume %row_count0 [range(%row_count0, 1, 1048576)] : index + %lhs_span = index.assume %lhs_span0 [range(%lhs_span0, 1, 1073741824)] : index + %rhs_span = index.assume %rhs_span0 [range(%rhs_span0, 1, 1073741824)] : index %output_width0 = index.add %lhs_width, %rhs_width : index %output_width = index.assume %output_width0 [range(%output_width0, 2, 131072)] : index %element_count0 = index.mul %output_width, %row_count : index @@ -45,9 +45,7 @@ template.def<@ggml.copy.concat_dim0.f32_body> device @ggml_copy_concat_dim0_f32_ %in_bounds = index.cmp ult, %linear, %element_count : index %lhs_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %lhs, %rhs, %output : buffer, buffer, buffer - %lhs_view = buffer.view %lhs_noalias[%base] : buffer -> view<1073741824xf32> - %rhs_view = buffer.view %rhs_noalias[%base] : buffer -> view<1073741824xf32> - %output_view = buffer.view %output_noalias[%base] : buffer -> view<1073741824xf32> + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%element_count]xf32> scf.if %in_bounds { %row = index.div %linear, %output_width : index @@ -56,18 +54,20 @@ template.def<@ggml.copy.concat_dim0.f32_body> device @ggml_copy_concat_dim0_f32_ %value = scf.if %from_lhs -> (f32) { %lhs_row = index.mul %row, %lhs_width : index %lhs_index0 = index.add %lhs_row, %column : index - %lhs_index = index.assume %lhs_index0 [range(%lhs_index0, 0, 1073741823)] : index - %loaded = view.load %lhs_view[%lhs_index] : view<1073741824xf32> -> f32 + %lhs_index, %lhs_bound = index.assume %lhs_index0, %lhs_span [lt(%lhs_index0, %lhs_span)] : index, index + %lhs_view = buffer.view %lhs_noalias[%base] : buffer -> view<[%lhs_bound]xf32> + %loaded = view.load %lhs_view[%lhs_index] : view<[%lhs_bound]xf32> -> f32 scf.yield %loaded : f32 } else { %rhs_column = index.sub %column, %lhs_width : index %rhs_row = index.mul %row, %rhs_width : index %rhs_index0 = index.add %rhs_row, %rhs_column : index - %rhs_index = index.assume %rhs_index0 [range(%rhs_index0, 0, 1073741823)] : index - %loaded = view.load %rhs_view[%rhs_index] : view<1073741824xf32> -> f32 + %rhs_index, %rhs_bound = index.assume %rhs_index0, %rhs_span [lt(%rhs_index0, %rhs_span)] : index, index + %rhs_view = buffer.view %rhs_noalias[%base] : buffer -> view<[%rhs_bound]xf32> + %loaded = view.load %rhs_view[%rhs_index] : view<[%rhs_bound]xf32> -> f32 scf.yield %loaded : f32 } - view.store %value, %output_view[%linear] : f32, view<1073741824xf32> + view.store %value, %output_view[%linear] : f32, view<[%element_count]xf32> } template.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/transpose_tile.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/transpose_tile.loom index d2d92be5fa0f..bac75d38b078 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/transpose_tile.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/transpose_tile.loom @@ -1,14 +1,9 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -config.decl @ggml.copy_transpose_f16.row_count : %value: index where [range(%value, 32, 32768), mul(%value, 32)] -config.decl @ggml.copy_transpose_f16.column_count : %value: index where [range(%value, 32, 32768), mul(%value, 32)] +template.decl @ggml.copy.transpose_tile.f16_body(%rows: index, %columns: index, %input: buffer, %output: buffer) -template.decl @ggml.copy.transpose_tile.f16_body(%input: buffer, %output: buffer) - -template.def<@ggml.copy.transpose_tile.f16_body> device @ggml_copy_transpose_tile_f16_body(%input: buffer, %output: buffer) { - %rows = config.get @ggml.copy_transpose_f16.row_count : index - %columns = config.get @ggml.copy_transpose_f16.column_count : index +template.def<@ggml.copy.transpose_tile.f16_body> device @ggml_copy_transpose_tile_f16_body(%rows: index, %columns: index, %input: buffer, %output: buffer) { %c0 = index.constant 0 : index %c4 = index.constant 4 : index %c8 = index.constant 8 : index @@ -23,16 +18,18 @@ template.def<@ggml.copy.transpose_tile.f16_body> device @ggml_copy_transpose_til %row = index.div %wi, %c8 : index %col_word = index.rem %wi, %c8 : index %col = index.mul %col_word, %c4 : index - %source_row = index.add %row_origin, %row : index + %source_row0 = index.add %row_origin, %row : index + %source_row, %bounded_rows = index.assume %source_row0, %rows [lt(%source_row0, %rows)] : index, index %source_col = index.add %col_origin, %col : index + %source_col_end0 = index.add %source_col, %c4 : index + %source_col_end, %bounded_columns0 = index.assume %source_col_end0, %columns [le(%source_col_end0, %columns)] : index, index %src, %dst = buffer.assume.noalias %input, %output : buffer, buffer - %input_view = buffer.view %src[%zero] : buffer -> view<[%rows]x[%columns]xf16> - %output_view = buffer.view %dst[%zero] : buffer -> view<[%columns]x[%rows]xf16> + %input_view = buffer.view %src[%zero] : buffer -> view<[%bounded_rows]x[%bounded_columns0]xf16> %shared = buffer.alloca align(16) %bytes : buffer %shared_view = buffer.view %shared[%zero] : buffer -> view<32x40xf16> %transposed = encoding.layout.strided [1, 40] : encoding %shared_transpose_view = buffer.view %shared[%zero] : buffer -> view<32x32xf16, %transposed> - %loaded = vector.load %input_view[%source_row, %source_col] : view<[%rows]x[%columns]xf16> -> vector<4xf16> + %loaded = vector.load %input_view[%source_row, %source_col] : view<[%bounded_rows]x[%bounded_columns0]xf16> -> vector<4xf16> vector.store %loaded, %shared_view[%row, %col] : vector<4xf16>, view<32x40xf16> kernel.barrier scope(workgroup) ordering(acq_rel) %add0 = index.constant 0 : index @@ -48,8 +45,12 @@ template.def<@ggml.copy.transpose_tile.f16_body> device @ggml_copy_transpose_til %read_col3 = index.add %col, %add3 : index %s3 = view.load %shared_transpose_view[%row, %read_col3] : view<32x32xf16, %transposed> -> f16 %swapped = vector.from_elements %s0, %s1, %s2, %s3 : vector<4xf16> - %target_row = index.add %col_origin, %row : index + %target_row0 = index.add %col_origin, %row : index + %target_row, %bounded_columns = index.assume %target_row0, %bounded_columns0 [lt(%target_row0, %bounded_columns0)] : index, index %target_col = index.add %row_origin, %col : index - vector.store %swapped, %output_view[%target_row, %target_col] : vector<4xf16>, view<[%columns]x[%rows]xf16> + %target_col_end0 = index.add %target_col, %c4 : index + %target_col_end, %bounded_output_rows = index.assume %target_col_end0, %bounded_rows [le(%target_col_end0, %bounded_rows)] : index, index + %bounded_output_view = buffer.view %dst[%zero] : buffer -> view<[%bounded_columns]x[%bounded_output_rows]xf16> + vector.store %swapped, %bounded_output_view[%target_row, %target_col] : vector<4xf16>, view<[%bounded_columns]x[%bounded_output_rows]xf16> template.return } 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 0ffe41275340..845551afa86c 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 @@ -90,7 +90,6 @@ func.def public inline @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired func.return %result0, %result1 : vector<16xf16>, vector<16xf16> } - func.def inline @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { %c4 = index.constant 4 : index %c8 = index.constant 8 : index @@ -877,7 +876,6 @@ func.def inline @ggml_q3k_f16_vector4(%weight: buffer, %row_byte_base: offset, % func.return %values : vector<4xf16> } - func.def inline @ggml_iq2s_grid_lookup_i32(%grid_index: i32, %word_index: i32) -> (i32) { %c0_i32 = scalar.constant 0 : i32 %c1_i32 = scalar.constant 1 : i32 @@ -1219,6 +1217,7 @@ func.def inline @ggml_iq2s_grid_lookup_i32(%grid_index: i32, %word_index: i32) - %word = scalar.ori %low, %high_shifted : i32 func.return %word : i32 } + func.def inline @ggml_iq2s_value_f32(%grid_word: i32, %shift: i32, %signs: i32, %sign_bit: i32, %scale: f32) -> (f32) { %c0_i32 = scalar.constant 0 : i32 %c255_i32 = scalar.constant 255 : i32 @@ -1329,6 +1328,701 @@ func.def inline @ggml_iq2s_f16_vector4(%weight: buffer, %row_byte_base: offset, %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> func.return %values : vector<4xf16> } + +// IQ2_S-only decoder backed by the canonical 1024-entry packed grid. +func.def inline @ggml_iq2s_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq2_block: index, %iq2_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c32_i32 = scalar.constant 32 : i32 + %c64_i32 = scalar.constant 64 : i32 + %c128_i32 = scalar.constant 128 : i32 + %c768_i32 = scalar.constant 768 : i32 + %c05_f32 = scalar.constant 0.5 : f32 + %c025_f32 = scalar.constant 0.25 : f32 + %block_bytes = index.constant 82 : offset + %qs_offset = index.constant 2 : offset + %qh_offset = index.constant 66 : offset + %scales_offset = index.constant 74 : offset + %block_byte_add = index.scale %iq2_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %qh_byte_base = index.add %block_byte_base, %qh_offset : offset + %scales_byte_base = index.add %block_byte_base, %scales_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<64xi8> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<8xi8> + %scales_view = buffer.view %weight[%scales_byte_base] : buffer -> view<8xi8> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<2048xi32> + %bounded_group = index.assume %iq2_group [range(%iq2_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 63)] : index + %packet_in_group0 = index.rem %bounded_packet, %c8 : index + %l0 = index.div %packet_in_group0, %c2 : index + %l = index.assume %l0 [range(%l0, 0, 3)] : index + %word_index = index.rem %packet_in_group0, %c2 : index + %qs_group_base = index.mul %bounded_group, %c4 : index + %qs_index = index.add %qs_group_base, %l : index + %signs_index = index.add %qs_index, %c32 : index + %scale_half = index.div %l, %c2 : index + %uses_high_scale = index.cmp eq, %scale_half, %c1 : index + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %qs_i8 = view.load %qs_view[%qs_index] : view<64xi8> -> i8 + %qh_i8 = view.load %qh_view[%bounded_group] : view<8xi8> -> i8 + %signs_i8 = view.load %qs_view[%signs_index] : view<64xi8> -> i8 + %scale_i8 = view.load %scales_view[%bounded_group] : view<8xi8> -> i8 + %qs = scalar.extui %qs_i8 : i8 to i32 + %qh = scalar.extui %qh_i8 : i8 to i32 + %signs = scalar.extui %signs_i8 : i8 to i32 + %scale_byte = scalar.extui %scale_i8 : i8 to i32 + %scale_low = scalar.andi %scale_byte, %c15_i32 : i32 + %scale_high_shift = scalar.shrui %scale_byte, %c4_i32 : i32 + %scale_high = scalar.andi %scale_high_shift, %c15_i32 : i32 + %scale_i32 = scf.select %uses_high_scale, %scale_high, %scale_low : i32 + %scale_base = scalar.uitofp %scale_i32 : i32 to f32 + %scale_plus = scalar.addf %scale_base, %c05_f32 : f32 + %scale_mul = scalar.mulf %d, %scale_plus : f32 + %scale = scalar.mulf %scale_mul, %c025_f32 : f32 + %l_i32 = index.cast %l : index to i32 + %two_l = scalar.muli %l_i32, %c2_i32 : i32 + %qh_shift = scalar.subi %c8_i32, %two_l : i32 + %qh_shifted = scalar.shli %qh, %qh_shift : i32 + %qh_masked = scalar.andi %qh_shifted, %c768_i32 : i32 + %grid_index_i32 = scalar.ori %qs, %qh_masked : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 1023)] : index + %grid_word_base = index.mul %grid_index, %c2 : index + %grid_word_index0 = index.add %grid_word_base, %word_index : index + %grid_word_index = index.assume %grid_word_index0 [range(%grid_word_index0, 0, 2047)] : index + %grid_word = view.load %grid_view[%grid_word_index] : view<2048xi32> -> i32 + %uses_word1 = index.cmp eq, %word_index, %c1 : index + %sign_bit0 = scf.select %uses_word1, %c16_i32, %c1_i32 : i32 + %sign_bit1 = scf.select %uses_word1, %c32_i32, %c2_i32 : i32 + %sign_bit2 = scf.select %uses_word1, %c64_i32, %c4_i32 : i32 + %sign_bit3 = scf.select %uses_word1, %c128_i32, %c8_i32 : i32 + %v0 = func.call @ggml_iq2s_value_f32(%grid_word, %c0_i32, %signs, %sign_bit0, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v1 = func.call @ggml_iq2s_value_f32(%grid_word, %c8_i32, %signs, %sign_bit1, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v2 = func.call @ggml_iq2s_value_f32(%grid_word, %c16_i32, %signs, %sign_bit2, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v3 = func.call @ggml_iq2s_value_f32(%grid_word, %c24_i32, %signs, %sign_bit3, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_iq2s_global_f16_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq2_block: index, %iq2_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_iq2s_global_f32_vector4(%grid, %weight, %row_byte_base, %iq2_block, %iq2_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @ggml_iq2xs_sign_mask(%sign_code: i32) -> (i32) { + %c1_i32 = scalar.constant 1 : i32 + %c7_i32 = scalar.constant 7 : i32 + %population = scalar.ctpopi %sign_code : i32 + %parity = scalar.andi %population, %c1_i32 : i32 + %high_bit = scalar.shli %parity, %c7_i32 : i32 + %sign_mask = scalar.ori %sign_code, %high_bit : i32 + func.return %sign_mask : i32 +} + +// IQ2_XXS decoder backed by the canonical 256-entry packed grid. +func.def inline @ggml_iq2xxs_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq2_block: index, %iq2_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c28_i32 = scalar.constant 28 : i32 + %c32_i32 = scalar.constant 32 : i32 + %c64_i32 = scalar.constant 64 : i32 + %c127_i32 = scalar.constant 127 : i32 + %c128_i32 = scalar.constant 128 : i32 + %c255_i32 = scalar.constant 255 : i32 + %c05_f32 = scalar.constant 0.5 : f32 + %c025_f32 = scalar.constant 0.25 : f32 + %block_bytes = index.constant 66 : offset + %qs_offset = index.constant 2 : offset + %bounded_group = index.assume %iq2_group [range(%iq2_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq2_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<16xi32> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<512xi32> + %metadata_index0 = index.mul %bounded_group, %c2 : index + %metadata_index = index.add %metadata_index0, %c1 : index + %indices_word = view.load %qs_view[%metadata_index0] : view<16xi32> -> i32 + %metadata_word = view.load %qs_view[%metadata_index] : view<16xi32> -> i32 + %l = index.div %bounded_packet, %c2 : index + %word_index = index.rem %bounded_packet, %c2 : index + %l_i32 = index.cast %l : index to i32 + %grid_shift = scalar.muli %l_i32, %c8_i32 : i32 + %grid_shifted = scalar.shrui %indices_word, %grid_shift : i32 + %grid_index_i32 = scalar.andi %grid_shifted, %c255_i32 : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 255)] : index + %grid_word_base = index.mul %grid_index, %c2 : index + %grid_word_index0 = index.add %grid_word_base, %word_index : index + %grid_word_index = index.assume %grid_word_index0 [range(%grid_word_index0, 0, 511)] : index + %grid_word = view.load %grid_view[%grid_word_index] : view<512xi32> -> i32 + %sign_shift = scalar.muli %l_i32, %c7_i32 : i32 + %sign_shifted = scalar.shrui %metadata_word, %sign_shift : i32 + %sign_code = scalar.andi %sign_shifted, %c127_i32 : i32 + %signs = func.call @ggml_iq2xs_sign_mask(%sign_code) : (i32) -> (i32) + %scale_code = scalar.shrui %metadata_word, %c28_i32 : i32 + %scale_base = scalar.uitofp %scale_code : i32 to f32 + %scale_plus = scalar.addf %scale_base, %c05_f32 : f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scale_mul = scalar.mulf %d, %scale_plus : f32 + %scale = scalar.mulf %scale_mul, %c025_f32 : f32 + %uses_word1 = index.cmp eq, %word_index, %c1 : index + %sign_bit0 = scf.select %uses_word1, %c16_i32, %c1_i32 : i32 + %sign_bit1 = scf.select %uses_word1, %c32_i32, %c2_i32 : i32 + %sign_bit2 = scf.select %uses_word1, %c64_i32, %c4_i32 : i32 + %sign_bit3 = scf.select %uses_word1, %c128_i32, %c8_i32 : i32 + %v0 = func.call @ggml_iq2s_value_f32(%grid_word, %c0_i32, %signs, %sign_bit0, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v1 = func.call @ggml_iq2s_value_f32(%grid_word, %c8_i32, %signs, %sign_bit1, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v2 = func.call @ggml_iq2s_value_f32(%grid_word, %c16_i32, %signs, %sign_bit2, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v3 = func.call @ggml_iq2s_value_f32(%grid_word, %c24_i32, %signs, %sign_bit3, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// IQ2_XS decoder backed by the canonical 512-entry packed grid. +func.def inline @ggml_iq2xs_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq2_block: index, %iq2_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c32_i32 = scalar.constant 32 : i32 + %c64_i32 = scalar.constant 64 : i32 + %c128_i32 = scalar.constant 128 : i32 + %c511_i32 = scalar.constant 511 : i32 + %c05_f32 = scalar.constant 0.5 : f32 + %c025_f32 = scalar.constant 0.25 : f32 + %block_bytes = index.constant 74 : offset + %qs_offset = index.constant 2 : offset + %scales_offset = index.constant 66 : offset + %bounded_group = index.assume %iq2_group [range(%iq2_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq2_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %scales_byte_base = index.add %block_byte_base, %scales_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<32xi16> + %scales_view = buffer.view %weight[%scales_byte_base] : buffer -> view<8xi8> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<1024xi32> + %l = index.div %bounded_packet, %c2 : index + %word_index = index.rem %bounded_packet, %c2 : index + %group_qs_base = index.mul %bounded_group, %c4 : index + %qs_index0 = index.add %group_qs_base, %l : index + %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 31)] : index + %q_i16 = view.load %qs_view[%qs_index] : view<32xi16> -> i16 + %q = scalar.extui %q_i16 : i16 to i32 + %grid_index_i32 = scalar.andi %q, %c511_i32 : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 511)] : index + %grid_word_base = index.mul %grid_index, %c2 : index + %grid_word_index0 = index.add %grid_word_base, %word_index : index + %grid_word_index = index.assume %grid_word_index0 [range(%grid_word_index0, 0, 1023)] : index + %grid_word = view.load %grid_view[%grid_word_index] : view<1024xi32> -> i32 + %sign_code = scalar.shrui %q, %c9_i32 : i32 + %signs = func.call @ggml_iq2xs_sign_mask(%sign_code) : (i32) -> (i32) + %scale_i8 = view.load %scales_view[%bounded_group] : view<8xi8> -> i8 + %scale_byte = scalar.extui %scale_i8 : i8 to i32 + %scale_half = index.div %l, %c2 : index + %scale_half_i32 = index.cast %scale_half : index to i32 + %scale_shift = scalar.muli %scale_half_i32, %c4_i32 : i32 + %scale_shifted = scalar.shrui %scale_byte, %scale_shift : i32 + %scale_code = scalar.andi %scale_shifted, %c15_i32 : i32 + %scale_base = scalar.uitofp %scale_code : i32 to f32 + %scale_plus = scalar.addf %scale_base, %c05_f32 : f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scale_mul = scalar.mulf %d, %scale_plus : f32 + %scale = scalar.mulf %scale_mul, %c025_f32 : f32 + %uses_word1 = index.cmp eq, %word_index, %c1 : index + %sign_bit0 = scf.select %uses_word1, %c16_i32, %c1_i32 : i32 + %sign_bit1 = scf.select %uses_word1, %c32_i32, %c2_i32 : i32 + %sign_bit2 = scf.select %uses_word1, %c64_i32, %c4_i32 : i32 + %sign_bit3 = scf.select %uses_word1, %c128_i32, %c8_i32 : i32 + %v0 = func.call @ggml_iq2s_value_f32(%grid_word, %c0_i32, %signs, %sign_bit0, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v1 = func.call @ggml_iq2s_value_f32(%grid_word, %c8_i32, %signs, %sign_bit1, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v2 = func.call @ggml_iq2s_value_f32(%grid_word, %c16_i32, %signs, %sign_bit2, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %v3 = func.call @ggml_iq2s_value_f32(%grid_word, %c24_i32, %signs, %sign_bit3, %scale) : (i32, i32, i32, i32, f32) -> (f32) + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// IQ3_XXS decoder backed by the canonical 256-entry packed grid. +func.def inline @ggml_iq3xxs_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq3_block: index, %iq3_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c6_i32 = scalar.constant 6 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c28_i32 = scalar.constant 28 : i32 + %c127_i32 = scalar.constant 127 : i32 + %c025_f32 = scalar.constant 0.25 : f32 + %block_bytes = index.constant 98 : offset + %qs_offset = index.constant 2 : offset + %aux_offset = index.constant 66 : offset + %bounded_group = index.assume %iq3_group [range(%iq3_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq3_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %aux_byte_base = index.add %block_byte_base, %aux_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<64xi8> + %aux_view = buffer.view %weight[%aux_byte_base] : buffer -> view<32xi8> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<256xi32> + %iqs_base = index.mul %bounded_group, %c8 : index + %iqs0 = index.add %iqs_base, %bounded_packet : index + %iqs = index.assume %iqs0 [range(%iqs0, 0, 63)] : index + %grid_index_i8 = view.load %qs_view[%iqs] : view<64xi8> -> i8 + %grid_index_i32 = scalar.extui %grid_index_i8 : i8 to i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 255)] : index + %grid_word = view.load %grid_view[%grid_index] : view<256xi32> -> i32 + %aux_base = index.mul %bounded_group, %c4 : index + %aux1_index = index.add %aux_base, %c1 : index + %aux2_index = index.add %aux_base, %c2 : index + %aux3_index = index.add %aux_base, %c4 : index + %aux3_index0 = index.sub %aux3_index, %c1 : index + %aux0_i8 = view.load %aux_view[%aux_base] : view<32xi8> -> i8 + %aux1_i8 = view.load %aux_view[%aux1_index] : view<32xi8> -> i8 + %aux2_i8 = view.load %aux_view[%aux2_index] : view<32xi8> -> i8 + %aux3_i8 = view.load %aux_view[%aux3_index0] : view<32xi8> -> i8 + %aux0 = scalar.extui %aux0_i8 : i8 to i32 + %aux1_0 = scalar.extui %aux1_i8 : i8 to i32 + %aux2_0 = scalar.extui %aux2_i8 : i8 to i32 + %aux3_0 = scalar.extui %aux3_i8 : i8 to i32 + %aux1 = scalar.shli %aux1_0, %c8_i32 : i32 + %aux2 = scalar.shli %aux2_0, %c16_i32 : i32 + %aux3 = scalar.shli %aux3_0, %c24_i32 : i32 + %aux01 = scalar.ori %aux0, %aux1 : i32 + %aux23 = scalar.ori %aux2, %aux3 : i32 + %aux = scalar.ori %aux01, %aux23 : i32 + %pair = index.div %bounded_packet, %c2 : index + %pair_i32 = index.cast %pair : index to i32 + %sign_shift = scalar.muli %pair_i32, %c7_i32 : i32 + %sign_index0 = scalar.shrui %aux, %sign_shift : i32 + %sign_index = scalar.andi %sign_index0, %c127_i32 : i32 + %p1 = scalar.shrui %sign_index, %c1_i32 : i32 + %p2 = scalar.shrui %sign_index, %c2_i32 : i32 + %p3 = scalar.shrui %sign_index, %c3_i32 : i32 + %p4 = scalar.shrui %sign_index, %c4_i32 : i32 + %p5 = scalar.shrui %sign_index, %c5_i32 : i32 + %p6 = scalar.shrui %sign_index, %c6_i32 : i32 + %p01 = scalar.addi %sign_index, %p1 : i32 + %p23 = scalar.addi %p2, %p3 : i32 + %p45 = scalar.addi %p4, %p5 : i32 + %p0123 = scalar.addi %p01, %p23 : i32 + %p456 = scalar.addi %p45, %p6 : i32 + %parity_sum = scalar.addi %p0123, %p456 : i32 + %parity = scalar.andi %parity_sum, %c1_i32 : i32 + %parity_bit = scalar.shli %parity, %c7_i32 : i32 + %signs = scalar.ori %sign_index, %parity_bit : i32 + %half = index.rem %bounded_packet, %c2 : index + %half_i32 = index.cast %half : index to i32 + %half_shift = scalar.muli %half_i32, %c4_i32 : i32 + %shifted_signs = scalar.shrui %signs, %half_shift : i32 + %scale0 = scalar.shrui %aux, %c28_i32 : i32 + %scale2 = scalar.muli %scale0, %c2_i32 : i32 + %scale_code = scalar.addi %scale2, %c1_i32 : i32 + %scale_f32 = scalar.uitofp %scale_code : i32 to f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %d_quarter = scalar.mulf %d, %c025_f32 : f32 + %db = scalar.mulf %d_quarter, %scale_f32 : f32 + %v0_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c0_i32, %shifted_signs, %c1_i32) : (i32, i32, i32, i32) -> (f32) + %v1_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c8_i32, %shifted_signs, %c2_i32) : (i32, i32, i32, i32) -> (f32) + %v2_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c16_i32, %shifted_signs, %c4_i32) : (i32, i32, i32, i32) -> (f32) + %v3_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c24_i32, %shifted_signs, %c8_i32) : (i32, i32, i32, i32) -> (f32) + %v0 = scalar.mulf %db, %v0_base : f32 + %v1 = scalar.mulf %db, %v1_base : f32 + %v2 = scalar.mulf %db, %v2_base : f32 + %v3 = scalar.mulf %db, %v3_base : f32 + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// IQ3_S decoder backed by the canonical 512-entry packed grid. +func.def inline @ggml_iq3s_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq3_block: index, %iq3_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2 = index.constant 2 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c8 = index.constant 8 : index + %c8_i32 = scalar.constant 8 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16 = index.constant 16 : index + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c256_i32 = scalar.constant 256 : i32 + %block_bytes = index.constant 110 : offset + %qs_offset = index.constant 2 : offset + %qh_offset = index.constant 66 : offset + %signs_offset = index.constant 74 : offset + %scales_offset = index.constant 106 : offset + %bounded_group = index.assume %iq3_group [range(%iq3_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq3_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %qh_byte_base = index.add %block_byte_base, %qh_offset : offset + %signs_byte_base = index.add %block_byte_base, %signs_offset : offset + %scales_byte_base = index.add %block_byte_base, %scales_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<64xi8> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<8xi8> + %signs_view = buffer.view %weight[%signs_byte_base] : buffer -> view<32xi8> + %scales_view = buffer.view %weight[%scales_byte_base] : buffer -> view<4xi8> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<512xi32> + %iqs_group_base = index.mul %bounded_group, %c8 : index + %iqs0 = index.add %iqs_group_base, %bounded_packet : index + %iqs = index.assume %iqs0 [range(%iqs0, 0, 63)] : index + %signs_index0 = index.div %iqs, %c2 : index + %signs_index = index.assume %signs_index0 [range(%signs_index0, 0, 31)] : index + %scale_index0 = index.div %iqs, %c16 : index + %scale_index = index.assume %scale_index0 [range(%scale_index0, 0, 3)] : index + %iqs_mod8_index = index.rem %iqs, %c8 : index + %high_shift_index = index.sub %c8, %iqs_mod8_index : index + %high_shift = index.cast %high_shift_index : index to i32 + %scale_shift_index = index.rem %bounded_group, %c2 : index + %scale_shift_base = index.cast %scale_shift_index : index to i32 + %scale_shift = scalar.muli %scale_shift_base, %c4_i32 : i32 + %sign_shift_index = index.rem %bounded_packet, %c2 : index + %sign_shift_base = index.cast %sign_shift_index : index to i32 + %sign_shift = scalar.muli %sign_shift_base, %c4_i32 : i32 + %q_i8 = view.load %qs_view[%iqs] : view<64xi8> -> i8 + %qh_i8 = view.load %qh_view[%bounded_group] : view<8xi8> -> i8 + %signs_i8 = view.load %signs_view[%signs_index] : view<32xi8> -> i8 + %scale_i8 = view.load %scales_view[%scale_index] : view<4xi8> -> i8 + %q_i32 = scalar.extui %q_i8 : i8 to i32 + %qh_i32 = scalar.extui %qh_i8 : i8 to i32 + %signs_i32 = scalar.extui %signs_i8 : i8 to i32 + %scale_i32 = scalar.extui %scale_i8 : i8 to i32 + %qh_shifted = scalar.shli %qh_i32, %high_shift : i32 + %qh_high = scalar.andi %qh_shifted, %c256_i32 : i32 + %grid_index_i32 = scalar.ori %q_i32, %qh_high : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 511)] : index + %grid_word = view.load %grid_view[%grid_index] : view<512xi32> -> i32 + %scale_shifted = scalar.shrui %scale_i32, %scale_shift : i32 + %local_scale = scalar.andi %scale_shifted, %c15_i32 : i32 + %local_scale2 = scalar.muli %local_scale, %c2_i32 : i32 + %scale_code = scalar.addi %local_scale2, %c1_i32 : i32 + %scale_f32 = scalar.uitofp %scale_code : i32 to f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %db = scalar.mulf %d, %scale_f32 : f32 + %shifted_signs = scalar.shrui %signs_i32, %sign_shift : i32 + %v0_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c0_i32, %shifted_signs, %c1_i32) : (i32, i32, i32, i32) -> (f32) + %v1_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c8_i32, %shifted_signs, %c2_i32) : (i32, i32, i32, i32) -> (f32) + %v2_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c16_i32, %shifted_signs, %c4_i32) : (i32, i32, i32, i32) -> (f32) + %v3_base = func.call @ggml_iq3s_signed_value_f32(%grid_word, %c24_i32, %shifted_signs, %c8_i32) : (i32, i32, i32, i32) -> (f32) + %v0 = scalar.mulf %db, %v0_base : f32 + %v1 = scalar.mulf %db, %v1_base : f32 + %v2 = scalar.mulf %db, %v2_base : f32 + %v3 = scalar.mulf %db, %v3_base : f32 + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_iq1s_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq1_block: index, %iq1_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c32768_i32 = scalar.constant 32768 : i32 + %delta_positive = scalar.constant 0.125 : f32 + %delta_negative = scalar.constant -0.125 : f32 + %block_bytes = index.constant 50 : offset + %qs_offset = index.constant 2 : offset + %qh_offset = index.constant 34 : offset + %bounded_group = index.assume %iq1_group [range(%iq1_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq1_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qs_byte_base = index.add %block_byte_base, %qs_offset : offset + %qh_byte_base = index.add %block_byte_base, %qh_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qs_view = buffer.view %weight[%qs_byte_base] : buffer -> view<32xi8> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<8xi16> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<16384xi8> + %entry_in_group = index.div %bounded_packet, %c2 : index + %qs_group_base = index.mul %bounded_group, %c4 : index + %qs_index0 = index.add %qs_group_base, %entry_in_group : index + %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 31)] : index + %qs_i8 = view.load %qs_view[%qs_index] : view<32xi8> -> i8 + %qs_i32 = scalar.extui %qs_i8 : i8 to i32 + %qh_i16 = view.load %qh_view[%bounded_group] : view<8xi16> -> i16 + %qh_i32 = scalar.extui %qh_i16 : i16 to i32 + %entry_i32 = index.cast %entry_in_group : index to i32 + %high_shift = scalar.muli %entry_i32, %c3_i32 : i32 + %qh_shifted = scalar.shrui %qh_i32, %high_shift : i32 + %grid_high = scalar.andi %qh_shifted, %c7_i32 : i32 + %grid_high_shifted = scalar.shli %grid_high, %c8_i32 : i32 + %grid_index_i32 = scalar.ori %qs_i32, %grid_high_shifted : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 2047)] : index + %grid_byte_base0 = index.mul %grid_index, %c8 : index + %packet_half = index.rem %bounded_packet, %c2 : index + %grid_packet_add = index.mul %packet_half, %c4 : index + %grid_byte_base1 = index.add %grid_byte_base0, %grid_packet_add : index + %grid_byte_base = index.assume %grid_byte_base1 [range(%grid_byte_base1, 0, 16380)] : index + %grid_values_i8 = vector.load %grid_view[%grid_byte_base] : view<16384xi8> -> vector<4xi8> + %grid_values = vector.sitofp %grid_values_i8 : vector<4xi8> to vector<4xf32> + %negative_mask = scalar.andi %qh_i32, %c32768_i32 : i32 + %uses_negative_delta = scalar.cmpi ne, %negative_mask, %c0_i32 : i32 + %delta = scf.select %uses_negative_delta, %delta_negative, %delta_positive : f32 + %delta_vector = vector.splat %delta : vector<4xf32> + %adjusted = vector.addf %grid_values, %delta_vector : vector<4xf32> + %scale_code0 = scalar.shrui %qh_i32, %c12_i32 : i32 + %scale_code = scalar.andi %scale_code0, %c7_i32 : i32 + %scale_twice = scalar.muli %scale_code, %c2_i32 : i32 + %scale_odd = scalar.addi %scale_twice, %c1_i32 : i32 + %scale_odd_f32 = scalar.uitofp %scale_odd : i32 to f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scale = scalar.mulf %d, %scale_odd_f32 : f32 + %scale_vector = vector.splat %scale : vector<4xf32> + %result = vector.mulf %adjusted, %scale_vector : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_iq1m_global_f32_vector4(%grid: buffer, %weight: buffer, %row_byte_base: offset, %iq1_block: index, %iq1_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c6_i32 = scalar.constant 6 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c61440_i32 = scalar.constant 61440 : i32 + %delta_positive = scalar.constant 0.125 : f32 + %delta_negative = scalar.constant -0.125 : f32 + %block_bytes = index.constant 56 : offset + %qh_offset = index.constant 32 : offset + %scales_offset = index.constant 48 : offset + %bounded_group = index.assume %iq1_group [range(%iq1_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq1_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_offset : offset + %scales_byte_base = index.add %block_byte_base, %scales_offset : offset + %qs_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi8> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi8> + %scales_view = buffer.view %weight[%scales_byte_base] : buffer -> view<4xi16> + %grid_view = buffer.view %grid[%c0_offset] : buffer -> view<16384xi8> + %entry_in_group = index.div %bounded_packet, %c2 : index + %qs_group_base = index.mul %bounded_group, %c4 : index + %qs_index0 = index.add %qs_group_base, %entry_in_group : index + %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 31)] : index + %qs_i8 = view.load %qs_view[%qs_index] : view<32xi8> -> i8 + %qs_i32 = scalar.extui %qs_i8 : i8 to i32 + %qh_group_base = index.mul %bounded_group, %c2 : index + %entry_pair = index.div %entry_in_group, %c2 : index + %qh_index0 = index.add %qh_group_base, %entry_pair : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %qh_i8 = view.load %qh_view[%qh_index] : view<16xi8> -> i8 + %qh_i32 = scalar.extui %qh_i8 : i8 to i32 + %entry_nibble = index.rem %entry_in_group, %c2 : index + %qh_shift_index = index.mul %entry_nibble, %c4 : index + %qh_shift = index.cast %qh_shift_index : index to i32 + %qh_nibble0 = scalar.shrui %qh_i32, %qh_shift : i32 + %qh_nibble = scalar.andi %qh_nibble0, %c15_i32 : i32 + %grid_high = scalar.andi %qh_nibble, %c7_i32 : i32 + %grid_high_shifted = scalar.shli %grid_high, %c8_i32 : i32 + %grid_index_i32 = scalar.ori %qs_i32, %grid_high_shifted : i32 + %grid_index0 = index.cast %grid_index_i32 : i32 to index + %grid_index = index.assume %grid_index0 [range(%grid_index0, 0, 2047)] : index + %grid_byte_base0 = index.mul %grid_index, %c8 : index + %packet_half = index.rem %bounded_packet, %c2 : index + %grid_packet_add = index.mul %packet_half, %c4 : index + %grid_byte_base1 = index.add %grid_byte_base0, %grid_packet_add : index + %grid_byte_base = index.assume %grid_byte_base1 [range(%grid_byte_base1, 0, 16380)] : index + %grid_values_i8 = vector.load %grid_view[%grid_byte_base] : view<16384xi8> -> vector<4xi8> + %grid_values = vector.sitofp %grid_values_i8 : vector<4xi8> to vector<4xf32> + %delta_mask = scalar.andi %qh_nibble, %c8_i32 : i32 + %uses_negative_delta = scalar.cmpi ne, %delta_mask, %c0_i32 : i32 + %delta = scf.select %uses_negative_delta, %delta_negative, %delta_positive : f32 + %delta_vector = vector.splat %delta : vector<4xf32> + %adjusted = vector.addf %grid_values, %delta_vector : vector<4xf32> + %scale0_i16 = view.load %scales_view[%c0] : view<4xi16> -> i16 + %scale1_i16 = view.load %scales_view[%c1] : view<4xi16> -> i16 + %scale2_i16 = view.load %scales_view[%c2] : view<4xi16> -> i16 + %scale3_i16 = view.load %scales_view[%c3] : view<4xi16> -> i16 + %scale0_i32 = scalar.extui %scale0_i16 : i16 to i32 + %scale1_i32 = scalar.extui %scale1_i16 : i16 to i32 + %scale2_i32 = scalar.extui %scale2_i16 : i16 to i32 + %scale3_i32 = scalar.extui %scale3_i16 : i16 to i32 + %base0_shifted = scalar.shrui %scale0_i32, %c12_i32 : i32 + %base0 = scalar.andi %base0_shifted, %c15_i32 : i32 + %base1_shifted = scalar.shrui %scale1_i32, %c12_i32 : i32 + %base1_masked = scalar.andi %base1_shifted, %c15_i32 : i32 + %base1 = scalar.shli %base1_masked, %c4_i32 : i32 + %base2_shifted = scalar.shrui %scale2_i32, %c12_i32 : i32 + %base2_masked = scalar.andi %base2_shifted, %c15_i32 : i32 + %base2 = scalar.shli %base2_masked, %c8_i32 : i32 + %base3_masked = scalar.andi %scale3_i32, %c61440_i32 : i32 + %base01 = scalar.ori %base0, %base1 : i32 + %base23 = scalar.ori %base2, %base3_masked : i32 + %base_bits = scalar.ori %base01, %base23 : i32 + %base_i16 = scalar.trunci %base_bits : i32 to i16 + %base_i16_vector = vector.splat %base_i16 : vector<1xi16> + %base_f16_vector = vector.bitcast %base_i16_vector : vector<1xi16> to vector<1xf16> + %base_f16 = vector.extract %base_f16_vector[0] : vector<1xf16> -> f16 + %base = scalar.extf %base_f16 : f16 to f32 + %scale_word_index = index.div %bounded_group, %c2 : index + %scale_word = view.load %scales_view[%scale_word_index] : view<4xi16> -> i16 + %scale_word_i32 = scalar.extui %scale_word : i16 to i32 + %group_half = index.rem %bounded_group, %c2 : index + %group_half_i32 = index.cast %group_half : index to i32 + %group_shift = scalar.muli %group_half_i32, %c6_i32 : i32 + %packet_scale_half = index.div %bounded_packet, %c4 : index + %packet_scale_i32 = index.cast %packet_scale_half : index to i32 + %packet_shift = scalar.muli %packet_scale_i32, %c3_i32 : i32 + %scale_shift = scalar.addi %group_shift, %packet_shift : i32 + %local_scale0 = scalar.shrui %scale_word_i32, %scale_shift : i32 + %local_scale = scalar.andi %local_scale0, %c7_i32 : i32 + %local_scale_twice = scalar.muli %local_scale, %c2_i32 : i32 + %local_scale_odd = scalar.addi %local_scale_twice, %c1_i32 : i32 + %local_scale_f32 = scalar.uitofp %local_scale_odd : i32 to f32 + %scale = scalar.mulf %base, %local_scale_f32 : f32 + %scale_vector = vector.splat %scale : vector<4xf32> + %result = vector.mulf %adjusted, %scale_vector : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def public inline @ggml_codebook_f32_vector4(%weight_format: index, %codebook: buffer, %weight: buffer, %row_byte_base: offset, %quant_block: index, %quant_group: index, %packet: index) -> (vector<4xf32>) { + %iq1_s_format = index.constant 19 : index + %iq3_xxs_format = index.constant 18 : index + %iq3_s_format = index.constant 21 : index + %iq1_m_format = index.constant 29 : index + %iq2_xxs_format = index.constant 24 : index + %iq2_xs_format = index.constant 25 : index + %is_iq1_s = index.cmp eq, %weight_format, %iq1_s_format : index + %is_iq3_xxs = index.cmp eq, %weight_format, %iq3_xxs_format : index + %is_iq3_s = index.cmp eq, %weight_format, %iq3_s_format : index + %is_iq1_m = index.cmp eq, %weight_format, %iq1_m_format : index + %is_iq2_xxs = index.cmp eq, %weight_format, %iq2_xxs_format : index + %is_iq2_xs = index.cmp eq, %weight_format, %iq2_xs_format : index + %other_values = scf.if %is_iq3_s -> (vector<4xf32>) { + %values = func.call @ggml_iq3s_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %iq2_values = scf.if %is_iq2_xxs -> (vector<4xf32>) { + %values = func.call @ggml_iq2xxs_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %xs_or_s = scf.if %is_iq2_xs -> (vector<4xf32>) { + %values = func.call @ggml_iq2xs_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %values = func.call @ggml_iq2s_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } + scf.yield %xs_or_s : vector<4xf32> + } + scf.yield %iq2_values : vector<4xf32> + } + %iq3_xxs_or_other = scf.if %is_iq3_xxs -> (vector<4xf32>) { + %values = func.call @ggml_iq3xxs_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %other_values : vector<4xf32> + } + %iq1_m_or_other = scf.if %is_iq1_m -> (vector<4xf32>) { + %values = func.call @ggml_iq1m_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %iq3_xxs_or_other : vector<4xf32> + } + %result = scf.if %is_iq1_s -> (vector<4xf32>) { + %values = func.call @ggml_iq1s_global_f32_vector4(%codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %iq1_m_or_other : vector<4xf32> + } + func.return %result : vector<4xf32> +} + +func.def public inline @ggml_codebook_f16_vector4(%weight_format: index, %codebook: buffer, %weight: buffer, %row_byte_base: offset, %quant_block: index, %quant_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_codebook_f32_vector4(%weight_format, %codebook, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + func.def inline @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>) { %grid_low0 = scalar.constant 257.0 : f32 %grid_low1 = scalar.constant 259.0 : f32 @@ -1634,40 +2328,34 @@ func.def inline @ggml_iq4nl_lookup_i8_vector4(%value_table: vector<16xi8>, %code } func.def inline @ggml_iq4nl_table_codes4(%q_bytes: vector<4xi8>, %uses_high: i1) -> (vector<4xi8>) { + %c0_i32 = scalar.constant 0 : i32 %c4_i32 = scalar.constant 4 : i32 %c8_i32 = scalar.constant 8 : i32 %c12_i32 = scalar.constant 12 : i32 %c16_i32 = scalar.constant 16 : i32 + %c20_i32 = scalar.constant 20 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c28_i32 = scalar.constant 28 : i32 %mask0_i32 = scalar.constant 15 : i32 - %mask1_i32 = scalar.constant 240 : i32 - %mask2_i32 = scalar.constant 3840 : i32 - %mask3_i32 = scalar.constant 61440 : i32 %q_word_vector = vector.bitcast %q_bytes : vector<4xi8> to vector<1xi32> %word = vector.extract %q_word_vector[0] : vector<1xi32> -> i32 - %word_shr4 = scalar.shrui %word, %c4_i32 : i32 - %word_shr8 = scalar.shrui %word, %c8_i32 : i32 - %word_shr12 = scalar.shrui %word, %c12_i32 : i32 - %word_shr16 = scalar.shrui %word, %c16_i32 : i32 - %low0 = scalar.andi %word, %mask0_i32 : i32 - %low1 = scalar.andi %word_shr4, %mask1_i32 : i32 - %low2 = scalar.andi %word_shr8, %mask2_i32 : i32 - %low3 = scalar.andi %word_shr12, %mask3_i32 : i32 - %low01 = scalar.ori %low0, %low1 : i32 - %low23 = scalar.ori %low2, %low3 : i32 - %low = scalar.ori %low01, %low23 : i32 - %high0 = scalar.andi %word_shr4, %mask0_i32 : i32 - %high1 = scalar.andi %word_shr8, %mask1_i32 : i32 - %high2 = scalar.andi %word_shr12, %mask2_i32 : i32 - %high3 = scalar.andi %word_shr16, %mask3_i32 : i32 - %high01 = scalar.ori %high0, %high1 : i32 - %high23 = scalar.ori %high2, %high3 : i32 - %high = scalar.ori %high01, %high23 : i32 - %selected = scf.select %uses_high, %high, %low : i32 - %selected_shr8 = scalar.shrui %selected, %c8_i32 : i32 - %packed0 = scalar.trunci %selected : i32 to i8 - %packed1 = scalar.trunci %selected_shr8 : i32 to i8 - %packed = vector.from_elements %packed0, %packed1 : vector<2xi8> - %codes = vector.bitunpacku<4> %packed : vector<2xi8> -> vector<4xi8> + %base_shift = scf.select %uses_high, %c4_i32, %c0_i32 : i32 + %shift1 = scf.select %uses_high, %c12_i32, %c8_i32 : i32 + %shift2 = scf.select %uses_high, %c20_i32, %c16_i32 : i32 + %shift3 = scf.select %uses_high, %c28_i32, %c24_i32 : i32 + %value0 = scalar.shrui %word, %base_shift : i32 + %value1 = scalar.shrui %word, %shift1 : i32 + %value2 = scalar.shrui %word, %shift2 : i32 + %value3 = scalar.shrui %word, %shift3 : i32 + %code0_i32 = scalar.andi %value0, %mask0_i32 : i32 + %code1_i32 = scalar.andi %value1, %mask0_i32 : i32 + %code2_i32 = scalar.andi %value2, %mask0_i32 : i32 + %code3_i32 = scalar.andi %value3, %mask0_i32 : i32 + %code0 = scalar.trunci %code0_i32 : i32 to i8 + %code1 = scalar.trunci %code1_i32 : i32 to i8 + %code2 = scalar.trunci %code2_i32 : i32 to i8 + %code3 = scalar.trunci %code3_i32 : i32 to i8 + %codes = vector.from_elements %code0, %code1, %code2, %code3 : vector<4xi8> func.return %codes : vector<4xi8> } @@ -1682,6 +2370,8 @@ func.def inline @ggml_iq4nl_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: bu %code_byte_base = index.add %block_byte_base, %code_offset : offset %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<16xi8> + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 %packet_page = index.rem %bounded_packet, %c4 : index %code_index0 = index.mul %packet_page, %c4 : index %code_index = index.assume %code_index0 [range(%code_index0, 0, 12), mul(%code_index0, 4)] : index @@ -1690,8 +2380,6 @@ func.def inline @ggml_iq4nl_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: bu %codes = func.call @ggml_iq4nl_table_codes4(%q_bytes, %uses_high) : (vector<4xi8>, i1) -> (vector<4xi8>) %lookup_i8 = func.call @ggml_iq4nl_lookup_i8_vector4(%iq4nl_table, %codes) : (vector<16xi8>, vector<4xi8>) -> (vector<4xi8>) %lookup = vector.sitofp %lookup_i8 : vector<4xi8> to vector<4xf32> - %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 %d_vector = vector.splat %d : vector<4xf32> %result = vector.mulf %d_vector, %lookup : vector<4xf32> func.return %result : vector<4xf32> @@ -1950,14 +2638,19 @@ func.def inline @ggml_f32_f32_vector4(%weight: buffer, %row_byte_base: offset, % 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 %c3 = index.constant 11 : index %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %c18 = index.constant 18 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index %c22 = index.constant 22 : index %c23 = index.constant 23 : index + %c24 = index.constant 24 : index + %c25 = index.constant 25 : index + %iq1_m_format = index.constant 29 : index %c16 = index.constant 16 : index %bf16_format = index.constant 30 : index %c32 = index.constant 32 : index @@ -1977,10 +2670,15 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %q4_1_tile_bytes = index.constant 160 : offset %q5_0_tile_bytes = index.constant 176 : offset %q5_1_tile_bytes = index.constant 192 : offset + %iq1_s_tile_bytes = index.constant 50 : offset + %iq1_m_tile_bytes = index.constant 56 : offset %iq2_s_tile_bytes = index.constant 82 : offset + %iq3_xxs_tile_bytes = index.constant 98 : offset %iq4_nl_tile_bytes = index.constant 144 : offset %iq3_s_tile_bytes = index.constant 110 : offset %iq4_xs_tile_bytes = index.constant 136 : offset + %iq2_xxs_tile_bytes = index.constant 66 : offset + %iq2_xs_tile_bytes = index.constant 74 : offset %q8_0_tile_bytes = index.constant 272 : offset %q8_1_tile_bytes = index.constant 288 : offset %f16_tile_bytes = index.constant 512 : offset @@ -1995,10 +2693,15 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %is_q4_1 = index.cmp eq, %weight_format, %q4_1_format : index %is_q5_0 = index.cmp eq, %weight_format, %q5_0_format : index %is_q5_1 = index.cmp eq, %weight_format, %q5_1_format : index + %is_iq1_s = index.cmp eq, %weight_format, %iq1_s_format : index + %is_iq1_m = index.cmp eq, %weight_format, %iq1_m_format : index %is_iq2_s = index.cmp eq, %weight_format, %c22 : index + %is_iq3_xxs = index.cmp eq, %weight_format, %c18 : index %is_iq4_nl = index.cmp eq, %weight_format, %c20 : index %is_iq3_s = index.cmp eq, %weight_format, %c21 : index %is_iq4_xs = index.cmp eq, %weight_format, %c23 : index + %is_iq2_xxs = index.cmp eq, %weight_format, %c24 : index + %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_f16 = index.cmp eq, %weight_format, %c16 : index @@ -2013,11 +2716,16 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %selected_q4_1 = scf.select %is_q4_1, %q4_1_tile_bytes, %selected_q4_0 : offset %selected_q5_0 = scf.select %is_q5_0, %q5_0_tile_bytes, %selected_q4_1 : offset %selected_q5_1 = scf.select %is_q5_1, %q5_1_tile_bytes, %selected_q5_0 : offset - %selected_iq2_s = scf.select %is_iq2_s, %iq2_s_tile_bytes, %selected_q5_1 : offset - %selected_iq4_nl = scf.select %is_iq4_nl, %iq4_nl_tile_bytes, %selected_iq2_s : offset + %selected_iq1_s = scf.select %is_iq1_s, %iq1_s_tile_bytes, %selected_q5_1 : offset + %selected_iq1_m = scf.select %is_iq1_m, %iq1_m_tile_bytes, %selected_iq1_s : offset + %selected_iq2_s = scf.select %is_iq2_s, %iq2_s_tile_bytes, %selected_iq1_m : offset + %selected_iq3_xxs = scf.select %is_iq3_xxs, %iq3_xxs_tile_bytes, %selected_iq2_s : offset + %selected_iq4_nl = scf.select %is_iq4_nl, %iq4_nl_tile_bytes, %selected_iq3_xxs : offset %selected_iq3_s = scf.select %is_iq3_s, %iq3_s_tile_bytes, %selected_iq4_nl : offset %selected_iq4_xs = scf.select %is_iq4_xs, %iq4_xs_tile_bytes, %selected_iq3_s : offset - %selected_q8_0 = scf.select %is_q8_0, %q8_0_tile_bytes, %selected_iq4_xs : offset + %selected_iq2_xxs = scf.select %is_iq2_xxs, %iq2_xxs_tile_bytes, %selected_iq4_xs : offset + %selected_iq2_xs = scf.select %is_iq2_xs, %iq2_xs_tile_bytes, %selected_iq2_xxs : offset + %selected_q8_0 = scf.select %is_q8_0, %q8_0_tile_bytes, %selected_iq2_xs : offset %selected_q8_1 = scf.select %is_q8_1, %q8_1_tile_bytes, %selected_q8_0 : offset %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 @@ -2025,16 +2733,21 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse func.return %selected : offset } -func.def inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) { +func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) { %q1_0_format = index.constant 10 : index + %iq1_s_format = index.constant 19 : index %c3 = index.constant 11 : index %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %c18 = index.constant 18 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index %c22 = index.constant 22 : index %c23 = index.constant 23 : index + %c24 = index.constant 24 : index + %c25 = index.constant 25 : index + %iq1_m_format = index.constant 29 : index %c16 = index.constant 16 : index %bf16_format = index.constant 30 : index %c32_format = index.constant 32 : index @@ -2066,13 +2779,23 @@ func.def inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_si %uses_k_tile5 = index.cmp eq, %weight_format, %c21 : index %uses_k_tile6 = index.cmp eq, %weight_format, %c23 : index %uses_k_tile7 = index.cmp eq, %weight_format, %c22 : index + %uses_k_tile8 = index.cmp eq, %weight_format, %c24 : index + %uses_k_tile9 = index.cmp eq, %weight_format, %c25 : index + %uses_k_tile10 = index.cmp eq, %weight_format, %c18 : index + %uses_k_tile11 = index.cmp eq, %weight_format, %iq1_s_format : index + %uses_k_tile12 = index.cmp eq, %weight_format, %iq1_m_format : index %uses_k_tile01 = scalar.ori %uses_k_tile0, %uses_k_tile1 : i1 %uses_k_tile23 = scalar.ori %uses_k_tile2, %uses_k_tile3 : i1 %uses_k_tile45 = scalar.ori %uses_k_tile4, %uses_k_tile5 : i1 %uses_k_tile67 = scalar.ori %uses_k_tile6, %uses_k_tile7 : i1 + %uses_k_tile89 = scalar.ori %uses_k_tile8, %uses_k_tile9 : i1 %uses_k_tile0123 = scalar.ori %uses_k_tile01, %uses_k_tile23 : i1 %uses_k_tile4567 = scalar.ori %uses_k_tile45, %uses_k_tile67 : i1 - %uses_k_tile = scalar.ori %uses_k_tile0123, %uses_k_tile4567 : i1 + %uses_k_tile0_7 = scalar.ori %uses_k_tile0123, %uses_k_tile4567 : i1 + %uses_k_tile8_10 = scalar.ori %uses_k_tile89, %uses_k_tile10 : i1 + %uses_k_tile11_12 = scalar.ori %uses_k_tile11, %uses_k_tile12 : i1 + %uses_k_tile0_10 = scalar.ori %uses_k_tile0_7, %uses_k_tile8_10 : i1 + %uses_k_tile = scalar.ori %uses_k_tile0_10, %uses_k_tile11_12 : i1 %tile_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) %tile_count = index.div %hidden_size, %c256 : index %tile_row_bytes = index.scale %tile_count, %tile_bytes : index, offset -> offset @@ -2121,7 +2844,7 @@ func.def inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_si // 10: Q1_0 // 11/4/5/6: Q3_K/Q4_K/Q5_K/Q6_K // 40/41/50/51: Q4_0/Q4_1/Q5_0/Q5_1 -// 22/20/21/23: IQ2_S/IQ4_NL/IQ3_S/IQ4_XS +// 22/20/21/23/24/25: IQ2_S/IQ4_NL/IQ3_S/IQ4_XS/IQ2_XXS/IQ2_XS // 80/81: Q8_0/Q8_1 // 16/30/32: F16/BF16/F32 func.def inline @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>) { 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 6b3fea6048de..71dd2cc24799 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 @@ -26,6 +26,9 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at // Normalize and publish the lane-owned packets. Query-tail rows never // participate in the output store. scf.if %lane_has_output { + %c4 = index.constant 4 : index + %lane_output_end0 = index.add %lane_output_channel, %c4 : index + %safe_lane_output_channel, %lane_output_end = index.assume %lane_output_channel, %lane_output_end0 [le(%lane_output_end0, %value_head_size)] : index, index %sum0_scalar = vector.extract %final_sum[0] : vector<4xf32> -> f32 %sum1_scalar = vector.extract %final_sum[1] : vector<4xf32> -> f32 %sum2_scalar = vector.extract %final_sum[2] : vector<4xf32> -> f32 @@ -55,7 +58,7 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at %gate_token_offset0 = index.mul %query_token0, %gate_stride_token : index %gate_head_offset0 = index.mul %query_head, %gate_stride_head : index %gate_head_base0 = index.add %gate_token_offset0, %gate_head_offset0 : index - %gate_index0_0 = index.add %gate_head_base0, %lane_output_channel : index + %gate_index0_0 = index.add %gate_head_base0, %safe_lane_output_channel : index %gate_index0 = index.assume %gate_index0_0 [range(%gate_index0_0, 0, 268435452)] : index %raw_gate0 = vector.load %gate_view[%gate_index0] : view<268435456xf32> -> vector<4xf32> %negative_gate0 = vector.negf %raw_gate0 : vector<4xf32> @@ -67,10 +70,10 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at } else { scf.yield %normalized0 : vector<4xf32> } - vector.store %published0, %output_view[%query_token0, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + vector.store %published0, %output_view[%query_token0, %query_head, %safe_lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> scf.if %publish_f16 { %half0 = vector.fptrunc %published0 : vector<4xf32> to vector<4xf16> - vector.store %half0, %f16_output_view[%query_token0, %query_head, %lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + vector.store %half0, %f16_output_view[%query_token0, %query_head, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> } } scf.if %query_valid1 { @@ -78,7 +81,7 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at %gate_token_offset1 = index.mul %query_token1, %gate_stride_token : index %gate_head_offset1 = index.mul %query_head, %gate_stride_head : index %gate_head_base1 = index.add %gate_token_offset1, %gate_head_offset1 : index - %gate_index1_0 = index.add %gate_head_base1, %lane_output_channel : index + %gate_index1_0 = index.add %gate_head_base1, %safe_lane_output_channel : index %gate_index1 = index.assume %gate_index1_0 [range(%gate_index1_0, 0, 268435452)] : index %raw_gate1 = vector.load %gate_view[%gate_index1] : view<268435456xf32> -> vector<4xf32> %negative_gate1 = vector.negf %raw_gate1 : vector<4xf32> @@ -90,10 +93,10 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at } else { scf.yield %normalized1 : vector<4xf32> } - vector.store %published1, %output_view[%query_token1, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + vector.store %published1, %output_view[%query_token1, %query_head, %safe_lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> scf.if %publish_f16 { %half1 = vector.fptrunc %published1 : vector<4xf32> to vector<4xf16> - vector.store %half1, %f16_output_view[%query_token1, %query_head, %lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + vector.store %half1, %f16_output_view[%query_token1, %query_head, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> } } scf.if %query_valid2 { @@ -101,7 +104,7 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at %gate_token_offset2 = index.mul %query_token2, %gate_stride_token : index %gate_head_offset2 = index.mul %query_head, %gate_stride_head : index %gate_head_base2 = index.add %gate_token_offset2, %gate_head_offset2 : index - %gate_index2_0 = index.add %gate_head_base2, %lane_output_channel : index + %gate_index2_0 = index.add %gate_head_base2, %safe_lane_output_channel : index %gate_index2 = index.assume %gate_index2_0 [range(%gate_index2_0, 0, 268435452)] : index %raw_gate2 = vector.load %gate_view[%gate_index2] : view<268435456xf32> -> vector<4xf32> %negative_gate2 = vector.negf %raw_gate2 : vector<4xf32> @@ -113,10 +116,10 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at } else { scf.yield %normalized2 : vector<4xf32> } - vector.store %published2, %output_view[%query_token2, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + vector.store %published2, %output_view[%query_token2, %query_head, %safe_lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> scf.if %publish_f16 { %half2 = vector.fptrunc %published2 : vector<4xf32> to vector<4xf16> - vector.store %half2, %f16_output_view[%query_token2, %query_head, %lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + vector.store %half2, %f16_output_view[%query_token2, %query_head, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> } } scf.if %query_valid3 { @@ -124,7 +127,7 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at %gate_token_offset3 = index.mul %query_token3, %gate_stride_token : index %gate_head_offset3 = index.mul %query_head, %gate_stride_head : index %gate_head_base3 = index.add %gate_token_offset3, %gate_head_offset3 : index - %gate_index3_0 = index.add %gate_head_base3, %lane_output_channel : index + %gate_index3_0 = index.add %gate_head_base3, %safe_lane_output_channel : index %gate_index3 = index.assume %gate_index3_0 [range(%gate_index3_0, 0, 268435452)] : index %raw_gate3 = vector.load %gate_view[%gate_index3] : view<268435456xf32> -> vector<4xf32> %negative_gate3 = vector.negf %raw_gate3 : vector<4xf32> @@ -136,10 +139,10 @@ template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_at } else { scf.yield %normalized3 : vector<4xf32> } - vector.store %published3, %output_view[%query_token3, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + vector.store %published3, %output_view[%query_token3, %query_head, %safe_lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> scf.if %publish_f16 { %half3 = vector.fptrunc %published3 : vector<4xf32> to vector<4xf16> - vector.store %half3, %f16_output_view[%query_token3, %query_head, %lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + vector.store %half3, %f16_output_view[%query_token3, %query_head, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> } } } 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 30e578b32a58..24449ad70d5e 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 @@ -84,6 +84,7 @@ template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flas } 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 %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 @@ -173,10 +174,10 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %query_token1 = index.add %query_origin, %query_row1 : index %query_token2 = index.add %query_origin, %query_row2 : index %query_token3 = index.add %query_origin, %query_row3 : index - %query_valid0 = index.cmp ult, %query_token0, %query_token_count : index - %query_valid1 = index.cmp ult, %query_token1, %query_token_count : index - %query_valid2 = index.cmp ult, %query_token2, %query_token_count : index - %query_valid3 = index.cmp ult, %query_token3, %query_token_count : index + %query_valid0 = index.cmp ult, %query_token0, %bounded_query_token_count : index + %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 %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 @@ -205,7 +206,7 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %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<[%query_token_count]x[%bounded_key_value_token_count]xf16> + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_query_token_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 @@ -218,7 +219,7 @@ 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>(%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) : (index, index, index, index, index, index, f32, buffer, buffer) 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>) { // Cache the complete 16x64 mask tile before QK. Causal masks leave future @@ -228,25 +229,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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%query_token0, %key_token] : view<[%bounded_query_token_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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%query_token1, %key_token] : view<[%bounded_query_token_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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%query_token2, %key_token] : view<[%bounded_query_token_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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %value = view.load %mask_view[%query_token3, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 scf.yield %value : f16 } else { scf.yield %c0_f16 : f16 @@ -332,7 +333,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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %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_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -342,7 +343,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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %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_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -352,7 +353,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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %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_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -362,7 +363,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<[%query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %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_f32 = scalar.extf %mask_f16 : f16 to f32 %score = scalar.addf %raw_score, %mask_f32 : f32 scf.yield %score : f32 @@ -401,7 +402,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, %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) : (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.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/projection_epilogue.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/projection_epilogue.loom index 258937ac528f..4d6a19003c00 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/projection_epilogue.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/projection_epilogue.loom @@ -3,14 +3,12 @@ template.decl @llm.gated_delta_net.projection_epilogue.apply(%alpha: f32, %beta: f32, %bias: f32, %a_scale: f32, %zero: f32, %one: f32) -> (f32, f32) -template.decl @llm.gated_delta_net.projection_epilogue.standalone_body(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +template.decl @llm.gated_delta_net.projection_epilogue.standalone_body(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) template.decl @llm.gated_delta_net.publish_f32.projection_pair(%gate: f32, %beta: f32, %gate_output: view<1048576xf32>, %beta_output: view<1048576xf32>, %index: index) config.decl @llm.gated_delta_net.epilogue_head_count : %value: index where [range(%value, 1, 4096)] -config.decl @llm.gated_delta_net.epilogue_element_count : %value: index where [range(%value, 1, 1048576)] - config.decl @llm.gated_delta_net.epilogue_workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] template.def<@llm.gated_delta_net.projection_epilogue.apply> device @llm_gated_delta_net_projection_epilogue_apply(%alpha: f32, %beta: f32, %bias: f32, %a_scale: f32, %zero: f32, %one: f32) -> (f32, f32) { @@ -27,12 +25,11 @@ template.def<@llm.gated_delta_net.projection_epilogue.apply> device @llm_gated_d template.return %gate, %beta_result : f32, f32 } -template.def<@llm.gated_delta_net.projection_epilogue.standalone_body> device @llm_gated_delta_net_projection_epilogue_standalone_body(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) { +template.def<@llm.gated_delta_net.projection_epilogue.standalone_body> device @llm_gated_delta_net_projection_epilogue_standalone_body(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) { %base = index.constant 0 : offset - %elements0 = config.get @llm.gated_delta_net.epilogue_element_count : index %heads0 = config.get @llm.gated_delta_net.epilogue_head_count : index %workgroup_size = config.get @llm.gated_delta_net.epilogue_workgroup_size : index - %elements = index.assume %elements0 [range(%elements0, 1, 1048576)] : index + %elements = index.assume %element_count [range(%element_count, 1, 1048576)] : index %heads = index.assume %heads0 [range(%heads0, 1, 4096)] : index %workgroup0 = kernel.workgroup.id : index %workgroup = index.assume %workgroup0 [range(%workgroup0, 0, 32767)] : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q1_0_q8_1_x4_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q1_0_q8_1_x4_decode.loom new file mode 100644 index 000000000000..bc1840354ead --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q1_0_q8_1_x4_decode.loom @@ -0,0 +1,199 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) + +// Expands four adjacent native Q1_0 sign bits only at the integer dot boundary. +func.def inline @ggml_q1_0_sign_word(%code: i8, %bit_base: index) -> (vector<4xi8>) { + %nibble_mask = scalar.constant 15 : i32 + %spread_multiplier = scalar.constant 2113665 : i32 + %byte_lsb_mask = scalar.constant 16843009 : i32 + %byte_ones = scalar.constant 255 : i32 + %code_i32 = scalar.extui %code : i8 to i32 + %base_i32 = index.cast %bit_base : index to i32 + %shifted = scalar.shrui %code_i32, %base_i32 : i32 + %nibble = scalar.andi %shifted, %nibble_mask : i32 + %spread0 = scalar.muli %nibble, %spread_multiplier : i32 + %spread = scalar.andi %spread0, %byte_lsb_mask : i32 + %zero_lsb = scalar.xori %spread, %byte_lsb_mask : i32 + %negative_bytes = scalar.muli %zero_lsb, %byte_ones : i32 + %signed_i32 = scalar.ori %negative_bytes, %spread : i32 + %signed_word = vector.from_elements %signed_i32 : vector<1xi32> + %signed = vector.bitcast %signed_word : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +// Computes one wave32 lane's contribution for one native 128-value Q1_0 block. +func.def inline @ggml_q1_0_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q1_block: index, %lane: index) -> (f32) { + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 18 : offset + %code_byte_add = index.constant 2 : offset + %c0_i32v = vector.constant 0 : vector<1xi32> + %block_byte_add = index.scale %q1_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<16xi8> + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %code_index0 = index.div %bounded_lane, %c2 : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 15)] : index + %code = view.load %code_view[%code_index] : view<16xi8> -> i8 + %half_byte = index.rem %bounded_lane, %c2 : index + %bit_base = index.mul %half_byte, %c4 : index + %signed_weight = func.call @ggml_q1_0_sign_word(%code, %bit_base) : (i8, index) -> (vector<4xi8>) + %block_in_q1 = index.div %bounded_lane, %c8 : index + %word_in_block = index.rem %bounded_lane, %c8 : index + %q8_block_base = index.mul %q1_block, %c4 : index + %q8_block = index.add %q8_block_base, %block_in_q1 : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %word_in_block) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %dot = vector.dot4i %signed_weight, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dot_i32 = vector.extract %dot[0] : vector<1xi32> -> i32 + %dot_f32 = scalar.sitofp %dot_i32 : i32 to f32 + %weight_scaled = scalar.mulf %dot_f32, %d : f32 + %contribution = scalar.mulf %weight_scaled, %q8_d : f32 + func.return %contribution : f32 +} + +// Computes one wave32 lane's partial for a complete native Q1_0 row. +func.def inline @ggml_q1_0_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c128 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%accumulator = %c0_f32 : f32) -> (f32) { + %value = func.call @ggml_q1_0_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %accumulator, %value : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Reuses each Q8_1 activation packet across two adjacent Q1_0 output rows. +func.def inline @ggml_q1_0_q8_1_x4_block_rows2_lane(%weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q1_block: index, %lane: index) -> (f32, f32) { + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 18 : offset + %code_byte_add = index.constant 2 : offset + %c0_i32v = vector.constant 0 : vector<1xi32> + %block_byte_add = index.scale %q1_block, %block_bytes : index, offset -> offset + %row0_block_base = index.add %row0_byte_base, %block_byte_add : offset + %row1_block_base = index.add %row1_byte_base, %block_byte_add : offset + %row0_code_base = index.add %row0_block_base, %code_byte_add : offset + %row1_code_base = index.add %row1_block_base, %code_byte_add : offset + %row0_d_view = buffer.view %weight[%row0_block_base] : buffer -> view<1xf16> + %row1_d_view = buffer.view %weight[%row1_block_base] : buffer -> view<1xf16> + %row0_code_view = buffer.view %weight[%row0_code_base] : buffer -> view<16xi8> + %row1_code_view = buffer.view %weight[%row1_code_base] : buffer -> view<16xi8> + %row0_d_f16 = view.load %row0_d_view[0] : view<1xf16> -> f16 + %row1_d_f16 = view.load %row1_d_view[0] : view<1xf16> -> f16 + %row0_d = scalar.extf %row0_d_f16 : f16 to f32 + %row1_d = scalar.extf %row1_d_f16 : f16 to f32 + %code_index0 = index.div %bounded_lane, %c2 : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 15)] : index + %row0_code = view.load %row0_code_view[%code_index] : view<16xi8> -> i8 + %row1_code = view.load %row1_code_view[%code_index] : view<16xi8> -> i8 + %half_byte = index.rem %bounded_lane, %c2 : index + %bit_base = index.mul %half_byte, %c4 : index + %row0_sign = func.call @ggml_q1_0_sign_word(%row0_code, %bit_base) : (i8, index) -> (vector<4xi8>) + %row1_sign = func.call @ggml_q1_0_sign_word(%row1_code, %bit_base) : (i8, index) -> (vector<4xi8>) + %block_in_q1 = index.div %bounded_lane, %c8 : index + %word_in_block = index.rem %bounded_lane, %c8 : index + %q8_block_base = index.mul %q1_block, %c4 : index + %q8_block = index.add %q8_block_base, %block_in_q1 : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %word_in_block) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %row0_dot = vector.dot4i %row0_sign, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %row1_dot = vector.dot4i %row1_sign, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %row0_i32 = vector.extract %row0_dot[0] : vector<1xi32> -> i32 + %row1_i32 = vector.extract %row1_dot[0] : vector<1xi32> -> i32 + %row0_f32 = scalar.sitofp %row0_i32 : i32 to f32 + %row1_f32 = scalar.sitofp %row1_i32 : i32 to f32 + %row0_scaled = scalar.mulf %row0_f32, %row0_d : f32 + %row1_scaled = scalar.mulf %row1_f32, %row1_d : f32 + %row0_contribution = scalar.mulf %row0_scaled, %q8_d : f32 + %row1_contribution = scalar.mulf %row1_scaled, %q8_d : f32 + func.return %row0_contribution, %row1_contribution : f32, f32 +} + +func.def inline @ggml_q1_0_q8_1_x4_rows2_lane(%input_size: index, %weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c128 : index + %result0, %result1 = scf.for %block = [%c0 to %block_count step %c1](%acc0 = %c0_f32 : f32, %acc1 = %c0_f32 : f32) -> (f32, f32) unroll(%c4) { + %value0, %value1 = func.call @ggml_q1_0_q8_1_x4_block_rows2_lane(%weight, %row0_byte_base, %row1_byte_base, %q8_input, %q8_row_byte_base, %block, %lane) : (buffer, offset, offset, buffer, offset, index, index) -> (f32, f32) + %next0 = scalar.addf %acc0, %value0 : f32 + %next1 = scalar.addf %acc1, %value1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + func.return %result0, %result1 : f32, f32 +} + +// Reuses each Q8_1 activation packet across the gate and up projection rows. +func.def inline @ggml_q1_0_q8_1_x4_pair_lane(%input_size: index, %gate: buffer, %gate_row_byte_base: offset, %up: buffer, %up_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : 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 + %c128 = index.constant 128 : index + %block_bytes = index.constant 18 : offset + %code_byte_add = index.constant 2 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_i32v = vector.constant 0 : vector<1xi32> + %block_count = index.div %bounded_input_size, %c128 : index + %gate_result, %up_result = scf.for %block = [%c0 to %block_count step %c1](%gate_acc = %c0_f32 : f32, %up_acc = %c0_f32 : f32) -> (f32, f32) unroll(%c4) { + %block_byte_add = index.scale %block, %block_bytes : index, offset -> offset + %gate_block_base = index.add %gate_row_byte_base, %block_byte_add : offset + %up_block_base = index.add %up_row_byte_base, %block_byte_add : offset + %gate_code_base = index.add %gate_block_base, %code_byte_add : offset + %up_code_base = index.add %up_block_base, %code_byte_add : offset + %gate_d_view = buffer.view %gate[%gate_block_base] : buffer -> view<1xf16> + %up_d_view = buffer.view %up[%up_block_base] : buffer -> view<1xf16> + %gate_code_view = buffer.view %gate[%gate_code_base] : buffer -> view<16xi8> + %up_code_view = buffer.view %up[%up_code_base] : buffer -> view<16xi8> + %gate_d_f16 = view.load %gate_d_view[0] : view<1xf16> -> f16 + %up_d_f16 = view.load %up_d_view[0] : view<1xf16> -> f16 + %gate_d = scalar.extf %gate_d_f16 : f16 to f32 + %up_d = scalar.extf %up_d_f16 : f16 to f32 + %code_index0 = index.div %bounded_lane, %c2 : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 15)] : index + %gate_code = view.load %gate_code_view[%code_index] : view<16xi8> -> i8 + %up_code = view.load %up_code_view[%code_index] : view<16xi8> -> i8 + %half_byte = index.rem %bounded_lane, %c2 : index + %bit_base = index.mul %half_byte, %c4 : index + %gate_sign = func.call @ggml_q1_0_sign_word(%gate_code, %bit_base) : (i8, index) -> (vector<4xi8>) + %up_sign = func.call @ggml_q1_0_sign_word(%up_code, %bit_base) : (i8, index) -> (vector<4xi8>) + %block_in_q1 = index.div %bounded_lane, %c8 : index + %word_in_block = index.rem %bounded_lane, %c8 : index + %q8_block_base = index.mul %block, %c4 : index + %q8_block = index.add %q8_block_base, %block_in_q1 : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %word_in_block) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %gate_dot = vector.dot4i %gate_sign, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %up_dot = vector.dot4i %up_sign, %q8_values, %c0_i32v : vector<4xi8>, vector<4xi8>, vector<1xi32> + %gate_i32 = vector.extract %gate_dot[0] : vector<1xi32> -> i32 + %up_i32 = vector.extract %up_dot[0] : vector<1xi32> -> i32 + %gate_f32 = scalar.sitofp %gate_i32 : i32 to f32 + %up_f32 = scalar.sitofp %up_i32 : i32 to f32 + %gate_scaled = scalar.mulf %gate_f32, %gate_d : f32 + %up_scaled = scalar.mulf %up_f32, %up_d : f32 + %gate_value = scalar.mulf %gate_scaled, %q8_d : f32 + %up_value = scalar.mulf %up_scaled, %q8_d : f32 + %next_gate = scalar.addf %gate_acc, %gate_value : f32 + %next_up = scalar.addf %up_acc, %up_value : f32 + scf.yield %next_gate, %next_up : f32, f32 + } + func.return %gate_result, %up_result : f32, f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom index 1ad784baf5f7..c7bf51776f16 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom @@ -8,7 +8,7 @@ func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %inpu func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) -config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 2, 5)] +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] // Shared contiguous 2D F32 activation skinny matmul pipeline. // diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..6012e42eef78 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_codebook_dequant_accumulate_f32.loom @@ -0,0 +1,215 @@ +config.def @ggml.mul_mat.activation_format = 0 : index + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Shared contiguous 2D F32 activation tiled matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_codebook_f16_vector4(%weight_format: index, %codebook: buffer, %weight: buffer, %row_byte_base: offset, %quant_block: index, %quant_group: index, %packet: index) -> (vector<4xf16>) + +func.def public inline @ggml_matmul_tiled_input_f32_codebook_dequant_accumulate_f32(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %bounded_token_count, %c31 : index + %token_tile_count = index.div %padded_token_count, %c32 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %bounded_input_size) : (index, index) -> (offset) + %split_count = kernel.workgroup.count : index + %partition = kernel.workgroup.id : index + %blocks_per_partition = index.div %quant_block_count, %split_count : index + %quant_block_begin = index.mul %partition, %blocks_per_partition : index + %quant_block_end = index.add %quant_block_begin, %blocks_per_partition : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f16_format = index.constant 1 : index + %uses_f16_activation = index.cmp eq, %activation_format, %f16_format : index + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %input_f32_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %input_f16_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf16> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_codebook_f16_vector4(%weight_format, %codebook, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %converted = scf.if %uses_f16_activation -> (vector<4xf16>) { + %loaded = vector.load.mask %input_f16_view[%bounded_token, %input_k], %mask, %c0_f16x4 : view<[%bounded_token_count]x[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + scf.yield %loaded : vector<4xf16> + } else { + %loaded = vector.load.mask %input_f32_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %narrow = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %narrow : vector<4xf16> + } + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile>(%bounded_token_count, %bounded_output_size, %output_tile_count, %token_tile_count, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + + +template.def<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core> device @ggml_matmul_tiled_input_f32_codebook_dequant_accumulate_f32_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + func.call @ggml_matmul_tiled_input_f32_codebook_dequant_accumulate_f32(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %codebook, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + 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 57a9543eb066..a4b588739396 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 @@ -15,7 +15,7 @@ config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1 // Op wrappers provide publication and tile-finalization hooks by defining: // ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 // ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile -template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) @@ -31,7 +31,7 @@ func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vecto func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) -func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { +func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index @@ -52,6 +52,8 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_form %c255 = index.constant 255 : index %c256 = index.constant 256 : index %c0_offset = index.constant 0 : offset + %f16_element_bytes = index.constant 2 : offset + %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 @@ -78,8 +80,6 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_form %f16_format = index.constant 1 : index %uses_f16_activation = index.cmp eq, %activation_format, %f16_format : index %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer - %input_f32_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> - %input_f16_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf16> %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer @@ -130,18 +130,23 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_form scf.if %is_activation_row { %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index %token = index.add %token_tile_base, %activation_row : index - %valid_token = index.cmp ult, %token, %bounded_token_count : index + %tail_valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_token = scalar.ori %full_token_tiles, %tail_valid_token : i1 %input_k = index.add %k_origin, %load_k : index %valid_input = index.cmp ult, %input_k, %bounded_input_size : index %valid_activation = scalar.andi %valid_token, %valid_input : i1 %activation_values = scf.if %valid_activation -> (vector<4xf16>) { - %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %row_element_offset = index.mul %token, %bounded_input_size : index + %f16_row_byte_offset = index.scale %row_element_offset, %f16_element_bytes : index, offset -> offset + %f32_row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %input_f16_row = buffer.view %input_noalias[%f16_row_byte_offset] : buffer -> view<[%bounded_input_size]xf16> + %input_f32_row = buffer.view %input_noalias[%f32_row_byte_offset] : buffer -> view<[%bounded_input_size]xf32> %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> %converted = scf.if %uses_f16_activation -> (vector<4xf16>) { - %loaded = vector.load.mask %input_f16_view[%bounded_token, %input_k], %mask, %c0_f16x4 : view<[%bounded_token_count]x[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %loaded = vector.load.mask %input_f16_row[%input_k], %mask, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> scf.yield %loaded : vector<4xf16> } else { - %loaded = vector.load.mask %input_f32_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %loaded = vector.load.mask %input_f32_row[%input_k], %mask, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> %narrow = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> scf.yield %narrow : vector<4xf16> } @@ -181,8 +186,10 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_form %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index %channel1_base = index.add %subgroup_channel_base, %c16 : index %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index - %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %tail_valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %tail_valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_token0 = scalar.ori %full_token_tiles, %tail_valid_token0 : i1 + %valid_token1 = scalar.ori %full_token_tiles, %tail_valid_token1 : i1 %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 @@ -224,9 +231,7 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_form func.return } - - -template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - func.call @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + func.call @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_tiles, %weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) template.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_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 new file mode 100644 index 000000000000..875207cae569 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32_token64.loom @@ -0,0 +1,240 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.mul_mat.activation_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Shared contiguous 2D F32 activation tiled matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32_token64.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c48 = index.constant 48 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %f16_element_bytes = index.constant 2 : offset + %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 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %bounded_token_count, %c63 : index + %token_tile_count = index.div %padded_token_count, %c64 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %bounded_input_size) : (index, index) -> (offset) + %split_count = kernel.workgroup.count : index + %partition = kernel.workgroup.id : index + %blocks_per_partition = index.div %quant_block_count, %split_count : index + %quant_block_begin = index.mul %partition, %blocks_per_partition : index + %quant_block_end = index.add %quant_block_begin, %blocks_per_partition : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f16_format = index.constant 1 : index + %uses_f16_activation = index.cmp eq, %activation_format, %f16_format : index + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %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> + %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 + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 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> + %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>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c32] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %activation_row = index.assume %local_row [range(%local_row, 0, 63)] : index + %token = index.add %token_tile_base, %activation_row : index + %tail_valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_token = scalar.ori %full_token_tiles, %tail_valid_token : i1 + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %row_element_offset = index.mul %token, %bounded_input_size : index + %f16_row_byte_offset = index.scale %row_element_offset, %f16_element_bytes : index, offset -> offset + %f32_row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %input_f16_row = buffer.view %input_noalias[%f16_row_byte_offset] : buffer -> view<[%bounded_input_size]xf16> + %input_f32_row = buffer.view %input_noalias[%f32_row_byte_offset] : buffer -> view<[%bounded_input_size]xf32> + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %converted = scf.if %uses_f16_activation -> (vector<4xf16>) { + %loaded = vector.load.mask %input_f16_row[%input_k], %mask, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + scf.yield %loaded : vector<4xf16> + } else { + %loaded = vector.load.mask %input_f32_row[%input_k], %mask, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %narrow = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %narrow : vector<4xf16> + } + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<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 { + %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> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next02, %next03 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result02, %block_result03 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %publish_token2 = index.add %publish_token, %c32 : index + %token2 = index.add %token_tile_base, %publish_token2 : index + %publish_token3 = index.add %publish_token, %c48 : index + %token3 = index.add %token_tile_base, %publish_token3 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %tail_valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %tail_valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %tail_valid_token2 = index.cmp ult, %token2, %bounded_token_count : index + %tail_valid_token3 = index.cmp ult, %token3, %bounded_token_count : index + %valid_token0 = scalar.ori %full_token_tiles, %tail_valid_token0 : i1 + %valid_token1 = scalar.ori %full_token_tiles, %tail_valid_token1 : i1 + %valid_token2 = scalar.ori %full_token_tiles, %tail_valid_token2 : i1 + %valid_token3 = scalar.ori %full_token_tiles, %tail_valid_token3 : i1 + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %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> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result02, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %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> + 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> + 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> + 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) + func.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32_token64.core> device @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64_core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + func.call @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_token64(%full_token_tiles, %weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom index 3dde5e2c456b..42fc4e44b4f7 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom @@ -1,7 +1,7 @@ // Shared paired tiled matmul pipeline for contiguous 2D F32 activations. // The core preserves the legacy generic paired tile schedule and delegates // pair postprocess and publication to the op wrapper. -template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) @@ -23,7 +23,7 @@ func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vecto func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) -func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { +func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_token_tiles: i1, %op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048)] : index %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index @@ -42,6 +42,7 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: in %c255 = index.constant 255 : index %c256 = index.constant 256 : index %c0_offset = index.constant 0 : offset + %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 @@ -57,7 +58,6 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: in %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) %input_noalias, %gate_weight_noalias, %up_weight_noalias = buffer.assume.noalias %input, %gate_weight, %up_weight : buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer @@ -121,14 +121,17 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: in scf.if %is_activation_row { %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index %token = index.add %token_tile_base, %activation_row : index - %valid_token = index.cmp ult, %token, %bounded_token_count : index + %tail_valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_token = scalar.ori %full_token_tiles, %tail_valid_token : i1 %input_k = index.add %k_origin, %load_k : index %valid_input = index.cmp ult, %input_k, %bounded_input_size : index %valid_activation = scalar.andi %valid_token, %valid_input : i1 %activation_values = scf.if %valid_activation -> (vector<4xf16>) { - %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %row_element_offset = index.mul %token, %bounded_input_size : index + %row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %input_row = buffer.view %input_noalias[%row_byte_offset] : buffer -> view<[%bounded_input_size]xf32> %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> - %loaded = vector.load.mask %input_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %loaded = vector.load.mask %input_row[%input_k], %mask, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> scf.yield %converted : vector<4xf16> } else { @@ -172,8 +175,10 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: in %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index %channel1_base = index.add %subgroup_channel_base, %c16 : index %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index - %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %tail_valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %tail_valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_token0 = scalar.ori %full_token_tiles, %tail_valid_token0 : i1 + %valid_token1 = scalar.ori %full_token_tiles, %tail_valid_token1 : i1 %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 @@ -226,7 +231,7 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op: in func.return } -template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32_core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { - func.call @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%op, %epilogue, %gate_weight_format, %up_weight_format, %token_count, %input_size0, %output_size0, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) +template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32_core(%full_token_tiles: i1, %op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { + func.call @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %token_count, %input_size0, %output_size0, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) template.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom index f2ebbfda17a6..bd5ec88c580a 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom @@ -16,16 +16,16 @@ template.def<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4> device @ggml_m } template.def<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4> device @ggml_matmul_tiled_postprocess_f32_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index %c1 = index.constant 1 : index - %c0_offset = index.constant 0 : offset + %f32_element_bytes = index.constant 4 : offset %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %residual_input_noalias = buffer.assume.noalias %residual_input : buffer - %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %row_element_offset = index.mul %token0, %output_size : index + %row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %residual_input_row = buffer.view %residual_input_noalias[%row_byte_offset] : buffer -> view<[%output_size]xf32> %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> - %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %residual = vector.load.mask %residual_input_row[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> %sum = vector.addf %values, %residual : vector<4xf32> template.return %sum : vector<4xf32> } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom index 2324766d6c61..b03b98130939 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom @@ -3,27 +3,29 @@ template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_ template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) template.def<@ggml.matmul.tiled_publish.store_vector4> device @ggml_matmul_tiled_publish_store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index %c1 = index.constant 1 : index %c16 = index.constant 16 : index %c32 = index.constant 32 : index - %c0_offset = index.constant 0 : offset + %f16_element_bytes = index.constant 2 : offset + %f32_element_bytes = index.constant 4 : offset %is_f16 = index.cmp eq, %format, %c16 : index %is_f32 = index.cmp eq, %format, %c32 : index %output_noalias = buffer.assume.noalias %output : buffer - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %row_element_offset = index.mul %token0, %output_size : index %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> %publish_f32 = scalar.andi %publish_word, %is_f32 : i1 scf.if %publish_f32 { - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> - vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + %row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %output_row = buffer.view %output_noalias[%row_byte_offset] : buffer -> view<[%output_size]xf32> + vector.store.mask %values, %output_row[%channel], %mask : vector<4xf32>, view<[%output_size]xf32>, vector<4xi1> } %publish_f16 = scalar.andi %publish_word, %is_f16 : i1 scf.if %publish_f16 { - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf16> + %row_byte_offset = index.scale %row_element_offset, %f16_element_bytes : index, offset -> offset + %output_row = buffer.view %output_noalias[%row_byte_offset] : buffer -> view<[%output_size]xf16> %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> - vector.store.mask %truncated, %output_view[%token, %channel], %mask : vector<4xf16>, view<[%token_count]x[%output_size]xf16>, vector<4xi1> + vector.store.mask %truncated, %output_row[%channel], %mask : vector<4xf16>, view<[%output_size]xf16>, vector<4xi1> } template.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_activation.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_activation.loom new file mode 100644 index 000000000000..3caa49390236 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_activation.loom @@ -0,0 +1,26 @@ +config.def @ggml.matmul.vector.activation_format = 0 : index + +config.def @ggml.matmul.vector.publish_format = 0 : index + +template.decl @ggml.matmul.vector_codebook_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +func.decl @ggml_matmul_vector_codebook_f32_decode_pair(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.launch> pure @ggml_matmul_vector_codebook_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c64 = index.constant 64 : index + %rounded_output = index.add %output_size, %c1 : index + %output_groups = index.div %rounded_output, %c2 : index + template.return %output_groups, %c1, %c1, %c64 : index, index, index, index +} + +template.def<@ggml.matmul.vector_codebook_activation.core> device @ggml_matmul_vector_codebook_activation_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %publish_format = config.get @ggml.matmul.vector.publish_format : index + func.call @ggml_matmul_vector_codebook_f32_decode_pair(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %publish_format, %epsilon, %channel_group, %token_group, %input, %weight, %codebook, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_decode_pair.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_decode_pair.loom new file mode 100644 index 000000000000..9fa8eb8c42f9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_codebook_decode_pair.loom @@ -0,0 +1,181 @@ +template.decl @ggml.matmul.vector_activation.publish_pair2(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: i1, %arg6: f32, %arg7: f32, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_codebook_f32_vector4(%weight_format: index, %codebook: buffer, %weight: buffer, %row_byte_base: offset, %quant_block: index, %quant_group: index, %packet: index) -> (vector<4xf32>) + +func.def inline @ggml_matmul_vector_decode_pair_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block0 = scf.select %valid_block, %block, %c0 : index + %safe_block, %launch_block_count = index.assume %safe_block0, %block_count [lt(%safe_block0, %block_count)] : index, index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32> + %vector_half_base = index.mul %vector_half, %c128 : index + %vector_offset = index.mul %vector_index, %c4 : index + %input_index0 = index.add %vector_half_base, %vector_offset : index + %input_index1 = index.add %input_index0, %c32 : index + %input_index2 = index.add %input_index0, %c64 : index + %input_index3 = index.add %input_index0, %c96 : index + %block_k_base = index.mul %safe_block, %c256 : index + %k0 = index.add %block_k_base, %input_index0 : index + %k1 = index.add %block_k_base, %input_index1 : index + %k2 = index.add %block_k_base, %input_index2 : index + %k3 = index.add %block_k_base, %input_index3 : index + %mask0 = vector.mask.range [%k0 to %input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %input_size step %c1] : index -> vector<4xi1> + %mask2 = vector.mask.range [%k2 to %input_size step %c1] : index -> vector<4xi1> + %mask3 = vector.mask.range [%k3 to %input_size step %c1] : index -> vector<4xi1> + %input0 = vector.load.mask %input_view[%token, %safe_block, %input_index0], %mask0, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input1 = vector.load.mask %input_view[%token, %safe_block, %input_index1], %mask1, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input2 = vector.load.mask %input_view[%token, %safe_block, %input_index2], %mask2, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input3 = vector.load.mask %input_view[%token, %safe_block, %input_index3], %mask3, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} +func.def inline @ggml_matmul_vector_decode_pair_block_row(%weight_format: index, %input_size: index, %row: index, %block: index, %lane: index, %codebook: buffer, %weight: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = vector.constant 0.0 : vector<4xf32> + %zero_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %packet = index.rem %itid, %c8 : index + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %group_base = index.mul %vector_half, %c4 : index + %group0 = index.add %group_base, %c0 : index + %group1 = index.add %group_base, %c1 : index + %group2 = index.add %group_base, %c2 : index + %group3 = index.add %group_base, %c3 : index + %block_k_base = index.mul %safe_block, %c256 : index + %half_base = index.mul %vector_half, %c128 : index + %packet_k = index.mul %packet, %c4 : index + %k_half = index.add %block_k_base, %half_base : index + %k0 = index.add %k_half, %packet_k : index + %k1 = index.add %k0, %c32 : index + %k2 = index.add %k0, %c64 : index + %k3 = index.add %k0, %c96 : index + %valid_k0 = index.cmp ult, %k0, %input_size : index + %valid_k1 = index.cmp ult, %k1, %input_size : index + %valid_k2 = index.cmp ult, %k2, %input_size : index + %valid_k3 = index.cmp ult, %k3, %input_size : index + %weight0 = scf.if %valid_k0 -> (vector<4xf32>) { + %decoded = func.call @ggml_codebook_f32_vector4(%weight_format, %codebook, %weight, %row_base, %safe_block, %group0, %packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight1 = scf.if %valid_k1 -> (vector<4xf32>) { + %decoded = func.call @ggml_codebook_f32_vector4(%weight_format, %codebook, %weight, %row_base, %safe_block, %group1, %packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight2 = scf.if %valid_k2 -> (vector<4xf32>) { + %decoded = func.call @ggml_codebook_f32_vector4(%weight_format, %codebook, %weight, %row_base, %safe_block, %group2, %packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight3 = scf.if %valid_k3 -> (vector<4xf32>) { + %decoded = func.call @ggml_codebook_f32_vector4(%weight_format, %codebook, %weight, %row_base, %safe_block, %group3, %packet) : (index, buffer, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %mul0 = vector.mulf %input0, %weight0 : vector<4xf32> + %mul1 = vector.mulf %input1, %weight1 : vector<4xf32> + %mul2 = vector.mulf %input2, %weight2 : vector<4xf32> + %mul3 = vector.mulf %input3, %weight3 : vector<4xf32> + %sum01 = vector.addf %mul0, %mul1 : vector<4xf32> + %sum23 = vector.addf %mul2, %mul3 : vector<4xf32> + %sum = vector.addf %sum01, %sum23 : vector<4xf32> + %block_contribution = vector.reduce %sum, %zero_f32 : vector<4xf32>, f32 + %contribution = scf.select %valid_block, %block_contribution, %zero_f32 : f32 + func.return %contribution : f32 +} + +func.def public inline @ggml_matmul_vector_codebook_f32_decode_pair(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %tokens = index.assume %token_count [range(%token_count, 1, 1)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %lane0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = scalar.constant 0.0 : f32 + %token_valid = index.cmp ult, %token_group, %tokens : index + %safe_token0 = scf.select %token_valid, %token_group, %c0 : index + %token, %launch_tokens = index.assume %safe_token0, %tokens [lt(%safe_token0, %tokens)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %channel_group, %c2 : index + %row_valid = index.cmp ult, %row00, %n : index + %safe_row0 = scf.select %row_valid, %row00, %c0 : index + %row0, %launch_n = index.assume %safe_row0, %n [lt(%safe_row0, %n)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_n : index + %publish_output = scalar.andi %token_valid, %row_valid : i1 + %padded_k = index.add %k, %c255 : index + %block_count = index.div %padded_k, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %zero : f32, %row_acc1 = %zero : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_matmul_vector_decode_pair_load_f32_block(%launch_tokens, %k, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_matmul_vector_decode_pair_block_row(%weight_format, %k, %row0, %block, %lane, %codebook, %weight_noalias, %input0, %input1, %input2, %input3) : (index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + %row1_contribution = func.call @ggml_matmul_vector_decode_pair_block_row(%weight_format, %k, %row1, %block, %lane, %codebook, %weight_noalias, %input0, %input1, %input2, %input3) : (index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %zero : f32 + } + %next0 = scalar.addf %row_acc0, %contribution0 : f32 + %next1 = scalar.addf %row_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %sum0 = kernel.workgroup.reduce %acc0 : f32 + %sum1 = kernel.workgroup.reduce %acc1 : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %publish_pair = scalar.andi %publish_output, %is_lane_zero : i1 + template.apply<@ggml.matmul.vector_activation.publish_pair2>(%publish_pair, %token, %row0, %launch_tokens, %launch_n, %row1_valid, %sum0, %sum1, %output_accumulation, %output_unary_op, %publish_format, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, i1, f32, f32, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + func.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom index b3e926e61b2d..8f12fb9d16e1 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom @@ -31,13 +31,13 @@ template.def<@ggml.rope_f32.rotate_scalar_pair> device @ggml_rope_f32_rotate_sca template.return %rotated_x, %rotated_y : f32, f32 } -template.decl @ggml.rope_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) +template.decl @ggml.rope_f32.body(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) template.decl @ggml.rope_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) template.decl @ggml.rope_token_scale_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) -template.decl @ggml.rope_f32.scaled_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) +template.decl @ggml.rope_f32.scaled_body(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) template.decl @ggml.rope_set_rows_f32.body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) @@ -151,18 +151,17 @@ template.def<@ggml.rope_f32.copy_packet> device @ggml_rope_f32_copy_packet(%inpu template.return %values : vector<4xf32> } -template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) { +template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output_scale: f32, %output: buffer) { %token_capacity = config.get @ggml.rope_f32.token_capacity : index %head_count0 = config.get @ggml.rope_f32.head_count : index %head_size0 = config.get @ggml.rope_f32.head_size : index %n_dims0 = config.get @ggml.rope_f32.n_dims : index %input_stride1_0 = config.get @ggml.rope_f32.input_stride1 : index %input_stride2_0 = config.get @ggml.rope_f32.input_stride2 : index - %input_span0 = config.get @ggml.rope_f32.input_span : index %mscale = config.get @ggml.rope_f32.mscale : f32 %mode = config.get @ggml.rope_f32.mode : index %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %bounded_input_span = index.assume %input_span [range(%input_span, 1, 1073741824)] : index %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index %head = kernel.workgroup.id : index @@ -234,15 +233,16 @@ template.def<@ggml.rope_f32.scaled_body> device @ggml_rope_f32_scaled_body(%toke template.return } -template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { +template.def<@ggml.rope_f32.body> device @ggml_rope_f32_body(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { %one = scalar.constant 1.0 : f32 - template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %one, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %input_span, %positions, %input, %theta, %freq_factors, %one, %output) : (index, index, buffer, buffer, buffer, buffer, f32, buffer) template.return } template.def<@ggml.rope_scale_f32.body> device @ggml_rope_scale_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { %output_scale = config.get @ggml.rope_scale_f32.scale : f32 - template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + %input_span = config.get @ggml.rope_f32.input_span : index + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %input_span, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, index, buffer, buffer, buffer, buffer, f32, buffer) template.return } @@ -256,7 +256,8 @@ template.def<@ggml.rope_token_scale_f32.body> device @ggml_rope_token_scale_f32_ %token = scf.select %valid_token, %token0, %c0 : index %scale_view = buffer.view %scales[%c0_offset] : buffer -> view<[%bounded_token_count]xf32> %output_scale = view.load %scale_view[%token] : view<[%bounded_token_count]xf32> -> f32 - template.apply<@ggml.rope_f32.scaled_body>(%token_count, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, buffer, buffer, buffer, buffer, f32, buffer) + %input_span = config.get @ggml.rope_f32.input_span : index + template.apply<@ggml.rope_f32.scaled_body>(%token_count, %input_span, %positions, %input, %theta, %freq_factors, %output_scale, %output) : (index, index, buffer, buffer, buffer, buffer, f32, buffer) template.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom index 6d68753088a9..df4defa17958 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom @@ -3,6 +3,10 @@ template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) +template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32_publish_q8(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer, %q8_output: buffer) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] config.decl @qwen3_moe.model.rms_epsilon : f32 @@ -102,3 +106,43 @@ template.def<@ggml.routing.weighted_reduce_next_rmsnorm_f32> device requires [#t } template.return } + +// Retains the F32 publication and packs the same normalized row for direct +// Q8 consumers. The reload keeps the legacy reduction and RMSNorm motif exact. +template.def<@ggml.routing.weighted_reduce_next_rmsnorm_f32_publish_q8> device requires [#target.subgroup.size<32>] @ggml_routing_weighted_reduce_next_rmsnorm_f32_publish_q8(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer, %q8_output: buffer) { + template.apply<@ggml.routing.weighted_reduce_next_rmsnorm_f32>(%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : (index, buffer, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + + %workitem = kernel.workitem.id : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c1024 = index.constant 1024 : index + %row_bytes = index.constant 2304 : offset + %scratch_bytes = index.constant 1152 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %c0_offset = index.constant 0 : offset + %channel_base = index.mul %workitem, %c4 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %active_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %active_token, %token0, %c0 : index + %token = index.assume %safe_token0 [lt(%safe_token0, %bounded_token_count)] : index + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %next_projection_input_noalias, %q8_output_noalias = buffer.assume.noalias %next_projection_input, %q8_output : buffer, buffer + %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x2048xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + %publish_word = scalar.constant true : i1 + + scf.if %active_token { + scf.for %stripe = [%c0 to %c2 step %c1] unroll { + %channel = index.madd %stripe, %c1024, %channel_base : index + %values = vector.load %next_projection_input_view[%token, %channel] : view<[%bounded_token_count]x2048xf32> -> vector<4xf32> + template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %q8_output_noalias) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/generic_apply.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/generic_apply.loom index adf6a0bab7f6..6fcf82c6e8ed 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/generic_apply.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/generic_apply.loom @@ -3,18 +3,14 @@ template.decl @ggml.unary_f32.apply(%op: index, %value: f32) -> (f32) -template.decl @llm.ssm_conv.generic_apply(%linear: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) +template.decl @llm.ssm_conv.generic_apply(%channel: index, %token: index, %sequence: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) -template.def<@llm.ssm_conv.generic_apply> device @llm_ssm_conv_generic_apply(%linear: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window_v: view<1073741824xf32>, %filter_v: view<1048576xf32>) -> (f32) { +template.def<@llm.ssm_conv.generic_apply> device @llm_ssm_conv_generic_apply(%channel: index, %token: index, %sequence: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window_v: view<1073741824xf32>, %filter_v: view<1048576xf32>) -> (f32) { %c0 = index.constant 0 : index %c1 = index.constant 1 : index %c0_f32 = scalar.constant 0.0 : f32 - %channel = index.rem %linear, %d_inner : index - %linear_div_inner = index.div %linear, %d_inner : index - %token = index.rem %linear_div_inner, %n_t : index - %sequence = index.div %linear_div_inner, %n_t : index - - %window_rows0 = index.add %n_t, %d_conv : index + %bounded_n_t = index.assume %n_t [range(%n_t, 1, 512)] : index + %window_rows0 = index.add %bounded_n_t, %d_conv : index %window_rows = index.sub %window_rows0, %c1 : index %window_sequence_span0 = index.mul %window_rows, %d_inner : index %window_sequence_span = index.assume %window_sequence_span0 [range(%window_sequence_span0, 32, 34603008), mul(%window_sequence_span0, 32)] : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_f32.loom index c37323b1fc14..db53e44317ff 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_f32.loom @@ -1,27 +1,26 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @ggml.copy.concat_dim0.f32_body(%lhs: buffer, %rhs: buffer, %output: buffer) +template.decl @ggml.copy.concat_dim0.f32_body(%row_count: index, %lhs_span: index, %rhs_span: index, %lhs: buffer, %rhs: buffer, %output: buffer) amdgpu.target @ggml_copy_f32_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.concat_dim0_f32.lhs_width : %value: index where [range(%value, 1, 65536)] config.decl @ggml.concat_dim0_f32.rhs_width : %value: index where [range(%value, 1, 65536)] -config.decl @ggml.concat_dim0_f32.row_count : %value: index where [range(%value, 1, 1048576)] -kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_concat_dim0_f32") @ggml_concat_dim0_f32() { +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_concat_dim0_f32") @ggml_concat_dim0_f32(%row_count: index, %lhs_span: index, %rhs_span: index) { %one = index.constant 1 : index %workgroup_size = index.constant 256 : index %rounding = index.constant 255 : index %lhs_width = config.get @ggml.concat_dim0_f32.lhs_width : index %rhs_width = config.get @ggml.concat_dim0_f32.rhs_width : index - %row_count = config.get @ggml.concat_dim0_f32.row_count : index %output_width = index.add %lhs_width, %rhs_width : index - %element_count = index.mul %output_width, %row_count : index + %element_count0 = index.mul %output_width, %row_count : index + %element_count = index.assume %element_count0 [range(%element_count0, 1, 1073741824)] : index %rounded = index.add %element_count, %rounding : index %workgroup_count = index.div %rounded, %workgroup_size : index kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%workgroup_size, %one, %one) : index -} launch(%lhs: buffer, %rhs: buffer, %output: buffer) { - template.apply<@ggml.copy.concat_dim0.f32_body>(%lhs, %rhs, %output) : (buffer, buffer, buffer) +} launch(%row_count: index, %lhs_span: index, %rhs_span: index, %lhs: buffer, %rhs: buffer, %output: buffer) where [range(%row_count, 1, 1048576), range(%lhs_span, 1, 1073741824), range(%rhs_span, 1, 1073741824)] { + template.apply<@ggml.copy.concat_dim0.f32_body>(%row_count, %lhs_span, %rhs_span, %lhs, %rhs, %output) : (index, index, index, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/transpose_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/transpose_f16.loom index 4d661d290d8c..3b24f8946c4a 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/transpose_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/transpose_f16.loom @@ -1,23 +1,22 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @ggml.copy.transpose_tile.f16_body(%input: buffer, %output: buffer) +template.decl @ggml.copy.transpose_tile.f16_body(%rows: index, %columns: index, %input: buffer, %output: buffer) amdgpu.target @ggml_copy_f32_gfx11_wave64 {subgroup_size = 64} -config.decl @ggml.copy_transpose_f16.row_count : %value: index where [range(%value, 32, 32768), mul(%value, 32)] -config.decl @ggml.copy_transpose_f16.column_count : %value: index where [range(%value, 32, 32768), mul(%value, 32)] - -kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_transpose_f16") @ggml_copy_transpose_f16() { - %rows = config.get @ggml.copy_transpose_f16.row_count : index - %columns = config.get @ggml.copy_transpose_f16.column_count : index +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_transpose_f16") @ggml_copy_transpose_f16(%row_count: index, %column_count: index) { + %rows = index.assume %row_count [range(%row_count, 32, 32768), mul(%row_count, 32)] : index + %columns = index.assume %column_count [range(%column_count, 32, 32768), mul(%column_count, 32)] : index %tile_size = index.constant 32 : index %one = index.constant 1 : index %x = index.div %columns, %tile_size : index %gy = index.div %rows, %tile_size : index %y = index.constant 256 : index kernel.launch.config workgroups(%x, %gy, %one) workgroup_size(%y, %one, %one) : index -} launch(%input: buffer, %output: buffer) { - template.apply<@ggml.copy.transpose_tile.f16_body>(%input, %output) : (buffer, buffer) +} launch(%row_count: index, %column_count: index, %input: buffer, %output: buffer) { + %rows = index.assume %row_count [range(%row_count, 32, 32768), mul(%row_count, 32)] : index + %columns = index.assume %column_count [range(%column_count, 32, 32768), mul(%column_count, 32)] : index + template.apply<@ggml.copy.transpose_tile.f16_body>(%rows, %columns, %input, %output) : (index, index, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_bc_publish_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_bc_publish_q8_1_x4.loom new file mode 100644 index 000000000000..b944ac8cb6f9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_bc_publish_q8_1_x4.loom @@ -0,0 +1,103 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +amdgpu.target @ggml_binary_bc_publish_q8_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.binary_bc_f32.op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.binary_bc_f32.src0_broadcast_dim0 : %value: index where [range(%value, 0, 1)] +config.decl @ggml.binary_bc_f32.src0_broadcast_dim1 : %value: index where [range(%value, 0, 1)] +config.decl @ggml.binary_bc_f32.src1_broadcast_dim0 : %value: index where [range(%value, 0, 1)] +config.decl @ggml.binary_bc_f32.src1_broadcast_dim1 : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_binary_bc_publish_q8_gfx11_wave32) export("ggml_binary_bc_f32_publish_q8_1_x4") @ggml_binary_bc_f32_publish_q8_1_x4(%token_count: index, %hidden_size: index, %src0_element_count: index, %src1_element_count: index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %hidden_size: index, %src0_element_count: index, %src1_element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 5), range(%hidden_size, 128, 32768), mul(%hidden_size, 128), range(%src0_element_count, 1, 134217728), range(%src1_element_count, 1, 134217728)] { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 5)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 128, 32768), mul(%hidden_size, 128)] : index + %lhs_count = index.assume %src0_element_count [range(%src0_element_count, 1, 134217728)] : index + %rhs_count = index.assume %src1_element_count [range(%src1_element_count, 1, 134217728)] : index + %op = config.get @ggml.binary_bc_f32.op : index + %lhs_bc0 = config.get @ggml.binary_bc_f32.src0_broadcast_dim0 : index + %lhs_bc1 = config.get @ggml.binary_bc_f32.src0_broadcast_dim1 : index + %rhs_bc0 = config.get @ggml.binary_bc_f32.src1_broadcast_dim0 : index + %rhs_bc1 = config.get @ggml.binary_bc_f32.src1_broadcast_dim1 : index + %token = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %lhs_bc0_set = index.cmp ne, %lhs_bc0, %c0 : index + %lhs_bc1_set = index.cmp ne, %lhs_bc1, %c0 : index + %rhs_bc0_set = index.cmp ne, %rhs_bc0, %c0 : index + %rhs_bc1_set = index.cmp ne, %rhs_bc1, %c0 : index + %lhs_row_width = scf.select %lhs_bc0_set, %c1, %bounded_hidden_size : index + %rhs_row_width = scf.select %rhs_bc0_set, %c1, %bounded_hidden_size : index + %lhs_token = scf.select %lhs_bc1_set, %c0, %token : index + %rhs_token = scf.select %rhs_bc1_set, %c0, %token : index + %lhs_token_base = index.mul %lhs_token, %lhs_row_width : index + %rhs_token_base = index.mul %rhs_token, %rhs_row_width : index + %word_element_add = index.mul %workitem, %c4 : index + %physical_group_count = index.div %bounded_hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %lhs_view = buffer.view %lhs[%c0_offset] : buffer -> view<[%lhs_count]xf32> + %rhs_view = buffer.view %rhs[%c0_offset] : buffer -> view<[%rhs_count]xf32> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + scf.for %channel = [%word_element_add to %bounded_hidden_size step %c1024] { + %channel1 = index.add %channel, %c1 : index + %channel2 = index.add %channel, %c2 : index + %channel3 = index.add %channel, %c3 : index + %lhs_channel0 = scf.select %lhs_bc0_set, %c0, %channel : index + %lhs_channel1 = scf.select %lhs_bc0_set, %c0, %channel1 : index + %lhs_channel2 = scf.select %lhs_bc0_set, %c0, %channel2 : index + %lhs_channel3 = scf.select %lhs_bc0_set, %c0, %channel3 : index + %rhs_channel0 = scf.select %rhs_bc0_set, %c0, %channel : index + %rhs_channel1 = scf.select %rhs_bc0_set, %c0, %channel1 : index + %rhs_channel2 = scf.select %rhs_bc0_set, %c0, %channel2 : index + %rhs_channel3 = scf.select %rhs_bc0_set, %c0, %channel3 : index + %lhs_index0 = index.add %lhs_token_base, %lhs_channel0 : index + %lhs_index1 = index.add %lhs_token_base, %lhs_channel1 : index + %lhs_index2 = index.add %lhs_token_base, %lhs_channel2 : index + %lhs_index3 = index.add %lhs_token_base, %lhs_channel3 : index + %rhs_index0 = index.add %rhs_token_base, %rhs_channel0 : index + %rhs_index1 = index.add %rhs_token_base, %rhs_channel1 : index + %rhs_index2 = index.add %rhs_token_base, %rhs_channel2 : index + %rhs_index3 = index.add %rhs_token_base, %rhs_channel3 : index + %lhs0 = view.load %lhs_view[%lhs_index0] : view<[%lhs_count]xf32> -> f32 + %lhs1 = view.load %lhs_view[%lhs_index1] : view<[%lhs_count]xf32> -> f32 + %lhs2 = view.load %lhs_view[%lhs_index2] : view<[%lhs_count]xf32> -> f32 + %lhs3 = view.load %lhs_view[%lhs_index3] : view<[%lhs_count]xf32> -> f32 + %rhs0 = view.load %rhs_view[%rhs_index0] : view<[%rhs_count]xf32> -> f32 + %rhs1 = view.load %rhs_view[%rhs_index1] : view<[%rhs_count]xf32> -> f32 + %rhs2 = view.load %rhs_view[%rhs_index2] : view<[%rhs_count]xf32> -> f32 + %rhs3 = view.load %rhs_view[%rhs_index3] : view<[%rhs_count]xf32> -> f32 + %result0 = template.apply<@ggml.binary_f32.apply>(%op, %lhs0, %rhs0) : (index, f32, f32) -> (f32) + %result1 = template.apply<@ggml.binary_f32.apply>(%op, %lhs1, %rhs1) : (index, f32, f32) -> (f32) + %result2 = template.apply<@ggml.binary_f32.apply>(%op, %lhs2, %rhs2) : (index, f32, f32) -> (f32) + %result3 = template.apply<@ggml.binary_f32.apply>(%op, %lhs3, %rhs3) : (index, f32, f32) -> (f32) + %result = vector.from_elements %result0, %result1, %result2, %result3 : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %publish_word = index.cmp ult, %channel, %bounded_hidden_size : index + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%publish_word, %token_output_byte_base, %channel, %result, %scratch_values, %scratch_d, %q8_output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_publish_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_publish_q8_1_x4.loom new file mode 100644 index 000000000000..214ed7940459 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_publish_q8_1_x4.loom @@ -0,0 +1,81 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +amdgpu.target @ggml_binary_publish_q8_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.binary_f32.op : %value: index where [range(%value, 0, 8)] +config.decl @ggml.binary_f32.ne0 : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.binary_f32.src0_stride1 : %value: index where [range(%value, 128, 1048576)] +config.decl @ggml.binary_f32.src1_stride1 : %value: index where [range(%value, 128, 1048576)] +config.decl @ggml.binary_f32.src0_span : %value: index where [range(%value, 128, 134217728)] +config.decl @ggml.binary_f32.src1_span : %value: index where [range(%value, 128, 134217728)] + +kernel.def target(@ggml_binary_publish_q8_gfx11_wave32) export("ggml_binary_f32_publish_q8_1_x4") @ggml_binary_f32_publish_q8_1_x4(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %lhs: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 5)] { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 5)] : index + %hidden_size0 = config.get @ggml.binary_f32.ne0 : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %lhs_stride = config.get @ggml.binary_f32.src0_stride1 : index + %rhs_stride = config.get @ggml.binary_f32.src1_stride1 : index + %lhs_count = config.get @ggml.binary_f32.src0_span : index + %rhs_count = config.get @ggml.binary_f32.src1_span : index + %op = config.get @ggml.binary_f32.op : index + %token = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %lhs_token_base = index.mul %token, %lhs_stride : index + %rhs_token_base = index.mul %token, %rhs_stride : index + %word_element_add = index.mul %workitem, %c4 : index + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %lhs_view = buffer.view %lhs[%c0_offset] : buffer -> view<[%lhs_count]xf32> + %rhs_view = buffer.view %rhs[%c0_offset] : buffer -> view<[%rhs_count]xf32> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + scf.for %channel = [%word_element_add to %hidden_size step %c1024] { + %lhs_base = index.add %lhs_token_base, %channel : index + %rhs_base = index.add %rhs_token_base, %channel : index + %lhs0 = view.load %lhs_view[%lhs_base] : view<[%lhs_count]xf32> -> f32 + %rhs0 = view.load %rhs_view[%rhs_base] : view<[%rhs_count]xf32> -> f32 + %lhs_index1 = index.add %lhs_base, %c1 : index + %rhs_index1 = index.add %rhs_base, %c1 : index + %lhs1 = view.load %lhs_view[%lhs_index1] : view<[%lhs_count]xf32> -> f32 + %rhs1 = view.load %rhs_view[%rhs_index1] : view<[%rhs_count]xf32> -> f32 + %c2 = index.constant 2 : index + %lhs_index2 = index.add %lhs_base, %c2 : index + %rhs_index2 = index.add %rhs_base, %c2 : index + %lhs2 = view.load %lhs_view[%lhs_index2] : view<[%lhs_count]xf32> -> f32 + %rhs2 = view.load %rhs_view[%rhs_index2] : view<[%rhs_count]xf32> -> f32 + %c3 = index.constant 3 : index + %lhs_index3 = index.add %lhs_base, %c3 : index + %rhs_index3 = index.add %rhs_base, %c3 : index + %lhs3 = view.load %lhs_view[%lhs_index3] : view<[%lhs_count]xf32> -> f32 + %rhs3 = view.load %rhs_view[%rhs_index3] : view<[%rhs_count]xf32> -> f32 + %result0 = template.apply<@ggml.binary_f32.apply>(%op, %lhs0, %rhs0) : (index, f32, f32) -> (f32) + %result1 = template.apply<@ggml.binary_f32.apply>(%op, %lhs1, %rhs1) : (index, f32, f32) -> (f32) + %result2 = template.apply<@ggml.binary_f32.apply>(%op, %lhs2, %rhs2) : (index, f32, f32) -> (f32) + %result3 = template.apply<@ggml.binary_f32.apply>(%op, %lhs3, %rhs3) : (index, f32, f32) -> (f32) + %result = vector.from_elements %result0, %result1, %result2, %result3 : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden_size]xf32> + %publish_word = index.cmp ult, %channel, %hidden_size : index + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%publish_word, %token_output_byte_base, %channel, %result, %scratch_values, %scratch_d, %q8_output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/projection_epilogue_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/projection_epilogue_f32.loom index afbe36a8a7b4..9cc814d4be1f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/projection_epilogue_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/projection_epilogue_f32.loom @@ -1,24 +1,23 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @llm.gated_delta_net.projection_epilogue.standalone_body(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +template.decl @llm.gated_delta_net.projection_epilogue.standalone_body(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) config.decl @llm.gated_delta_net.epilogue_head_count : %value: index where [range(%value, 1, 4096)] -config.decl @llm.gated_delta_net.epilogue_element_count : %value: index where [range(%value, 1, 1048576)] - config.decl @llm.gated_delta_net.epilogue_workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] -kernel.def export("llm_gated_delta_net_projection_epilogue_f32") @llm_gated_delta_net_projection_epilogue_f32() { +kernel.def export("llm_gated_delta_net_projection_epilogue_f32") @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) { %unit = index.constant 1 : index %cneg1 = index.constant -1 : index - %elements = config.get @llm.gated_delta_net.epilogue_element_count : index + %elements = index.assume %element_count [range(%element_count, 1, 1048576)] : index %workgroup_size = config.get @llm.gated_delta_net.epilogue_workgroup_size : index %rounding = index.add %workgroup_size, %cneg1 : index %rounded = index.add %elements, %rounding : index %workgroups = index.div %rounded, %workgroup_size : index kernel.launch.config workgroups(%workgroups, %unit, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index -} launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) { - template.apply<@llm.gated_delta_net.projection_epilogue.standalone_body>(%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : (buffer, buffer, buffer, buffer, buffer, buffer) +} launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) { + %elements = index.assume %element_count [range(%element_count, 1, 1048576)] : index + template.apply<@llm.gated_delta_net.projection_epilogue.standalone_body>(%elements, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : (index, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom index efea51923750..b4d982d0d6af 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom @@ -1,4 +1,4 @@ -template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) @@ -64,9 +64,10 @@ kernel.def target(@llm_attention_k_matmul_rope_set_rows_tiled_gfx11_wave64) @llm %c0 = index.constant 0 : index %epilogue = index.constant 0 : index %c0_f32 = scalar.constant 0.0 : f32 + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom index 958be7d63b5c..e70144aa4d02 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom @@ -1,4 +1,4 @@ -template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) @@ -58,9 +58,10 @@ kernel.def target(@llm_attention_q_matmul_rope_tiled_gfx11_wave64) @llm_attentio %c0 = index.constant 0 : index %epilogue = index.constant 0 : index %c0_f32 = scalar.constant 0.0 : f32 + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom index 2b72a2feb93b..129beb52a9f3 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom @@ -1,4 +1,4 @@ -template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) @@ -57,9 +57,10 @@ kernel.def target(@llm_attention_v_matmul_set_rows_tiled_gfx11_wave64) @llm_atte %c0 = index.constant 0 : index %epilogue = index.constant 0 : index %c0_f32 = scalar.constant 0.0 : f32 + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_q8_1_x4_packed_test.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_q8_1_x4_packed_test.loom new file mode 100644 index 000000000000..7859b6c5606c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_q8_1_x4_packed_test.loom @@ -0,0 +1,336 @@ +// Test-only direct comparison between the existing Q1_0 F32 decoder and the +// inactive native Q1_0 x Q8_1 x4 packed contraction motif. +func.decl @ggml_q1_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) + +func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) + +func.decl @ggml_q1_0_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +func.decl @ggml_q1_0_q8_1_x4_rows2_lane(%input_size: index, %weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) + +func.decl @ggml_q1_0_q8_1_x4_pair_lane(%input_size: index, %gate: buffer, %gate_row_byte_base: offset, %up: buffer, %up_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) + +amdgpu.target @ggml_q1_0_q8_1_x4_gfx11_wave32 {subgroup_size = 32} + +func.def inline @ggml_q1_0_q8_1_x4_generic_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c128 : index + %lane_k = index.mul %bounded_lane, %c4 : index + %block_in_q1 = index.div %bounded_lane, %c8 : index + %word_in_block = index.rem %bounded_lane, %c8 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%accumulator = %c0_f32 : f32) -> (f32) { + %block_k = index.mul %block, %c128 : index + %k = index.add %block_k, %lane_k : index + %weight_values = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %q8_block_base = index.mul %block, %c4 : index + %q8_block = index.add %q8_block_base, %block_in_q1 : index + %q8_values_i8, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %word_in_block) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %q8_values_f32 = vector.sitofp %q8_values_i8 : vector<4xi8> to vector<4xf32> + %q8_scale = vector.splat %q8_d : vector<4xf32> + %activation = vector.mulf %q8_values_f32, %q8_scale : vector<4xf32> + %products = vector.mulf %weight_values, %activation : vector<4xf32> + %contribution = vector.reduce %products, %c0_f32 : vector<4xf32>, f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +kernel.def target(@ggml_q1_0_q8_1_x4_gfx11_wave32) @ggml_q1_0_q8_1_x4_generic_control(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size2 = index.assume %row0, %bounded_output_size [lt(%row0, %bounded_output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %weight_block_bytes = index.constant 18 : offset + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_row_byte_base = index.constant 0 : offset + %lane_value = func.call @ggml_q1_0_q8_1_x4_generic_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %dot = kernel.subgroup.reduce %lane_value : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_output_size2]xf32> + view.store %dot, %output_view[%row] : f32, view<[%bounded_output_size2]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_q8_1_x4_gfx11_wave32) @ggml_q1_0_q8_1_x4_packed(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size2 = index.assume %row0, %bounded_output_size [lt(%row0, %bounded_output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %weight_block_bytes = index.constant 18 : offset + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_row_byte_base = index.constant 0 : offset + %lane_value = func.call @ggml_q1_0_q8_1_x4_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %dot = kernel.subgroup.reduce %lane_value : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_output_size2]xf32> + view.store %dot, %output_view[%row] : f32, view<[%bounded_output_size2]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_q8_1_x4_gfx11_wave32) @ggml_q1_0_q8_1_x4_packed_rows2(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c32 = index.constant 32 : index + %rounded = index.add %output_size, %c1 : index + %groups = index.div %rounded, %c2 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %group = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %weight_block_bytes = index.constant 18 : offset + %q8_row_byte_base = index.constant 0 : offset + %row00 = index.mul %group, %c2 : index + %row0, %bounded_output_size = index.assume %row00, %output_size [lt(%row00, %output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %has_row1 = index.cmp ult, %row1, %bounded_output_size : index + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %row0_base = index.scale %row0, %weight_row_bytes : index, offset -> offset + %partial0, %partial1 = scf.if %has_row1 -> (f32, f32) { + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + %row1_base = index.scale %bounded_row1, %weight_row_bytes : index, offset -> offset + %pair0, %pair1 = func.call @ggml_q1_0_q8_1_x4_rows2_lane(%bounded_input_size, %weight, %row0_base, %row1_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, offset, buffer, offset, index) -> (f32, f32) + scf.yield %pair0, %pair1 : f32, f32 + } else { + %single = func.call @ggml_q1_0_q8_1_x4_row_lane(%bounded_input_size, %weight, %row0_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %single, %c0_f32 : f32, f32 + } + %sum0 = kernel.subgroup.reduce %partial0 : f32 + %sum1 = kernel.subgroup.reduce %partial1 : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum0, %output_view[%row0] : f32, view<[%bounded_output_size]xf32> + scf.if %has_row1 { + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + view.store %sum1, %output_view[%bounded_row1] : f32, view<[%bounded_output_size]xf32> + } + } + kernel.return +} + +kernel.def target(@ggml_q1_0_q8_1_x4_gfx11_wave32) @ggml_q1_0_q8_1_x4_packed_pair(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %q8_input: buffer, %gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size2 = index.assume %row0, %bounded_output_size [lt(%row0, %bounded_output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %weight_block_bytes = index.constant 18 : offset + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_row_byte_base = index.constant 0 : offset + %gate_lane, %up_lane = func.call @ggml_q1_0_q8_1_x4_pair_lane(%bounded_input_size, %gate_weight, %weight_row_byte_base, %up_weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, buffer, offset, index) -> (f32, f32) + %gate = kernel.subgroup.reduce %gate_lane : f32 + %up = kernel.subgroup.reduce %up_lane : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %gate_view = buffer.view %gate_output[%q8_row_byte_base] : buffer -> view<[%bounded_output_size2]xf32> + %up_view = buffer.view %up_output[%q8_row_byte_base] : buffer -> view<[%bounded_output_size2]xf32> + view.store %gate, %gate_view[%row] : f32, view<[%bounded_output_size2]xf32> + view.store %up, %up_view[%row] : f32, view<[%bounded_output_size2]xf32> + } + kernel.return +} + +check.case public @ggml_q1_0_q8_1_x4_boundary_differential_case { + %input_size = check.literal value(128) : index + %output_size = check.literal value(3) : index + %q8_input = check.generate.fill value(61) : tensor<144xi8> + %weight = check.generate.fill value(62) : tensor<3x18xi8> + %generic_output = check.generate.fill value(0.0) : tensor<3xf32> + %packed_output = check.generate.fill value(1.0) : tensor<3xf32> + kernel.launch @ggml_q1_0_q8_1_x4_generic_control[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %generic_output) : [index, index](index, index, tensor<144xi8>, tensor<3x18xi8>, tensor<3xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %packed_output) : [index, index](index, index, tensor<144xi8>, tensor<3x18xi8>, tensor<3xf32>) + check.expect.close actual(%packed_output) expected(%generic_output) atol(0.01) rtol(1.0000000000000001e-05) nan(same) : tensor<3xf32> + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_bonsai_2560x4096_differential_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %generic_output = check.generate.fill value(0.0) : tensor<4096xf32> + %packed_output = check.generate.fill value(1.0) : tensor<4096xf32> + %rows2_output = check.generate.fill value(2.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_q8_1_x4_generic_control[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %generic_output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %packed_output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed_rows2[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %rows2_output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.expect.close actual(%packed_output) expected(%generic_output) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<4096xf32> + check.expect.close actual(%rows2_output) expected(%generic_output) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<4096xf32> + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_bonsai_2560x9728_pair_differential_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(9728) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %gate_weight = check.generate.fill value(62) : tensor<9728x20x18xi8> + %up_weight = check.generate.fill value(63) : tensor<9728x20x18xi8> + %generic_gate = check.generate.fill value(0.0) : tensor<9728xf32> + %generic_up = check.generate.fill value(0.0) : tensor<9728xf32> + %paired_gate = check.generate.fill value(1.0) : tensor<9728xf32> + %paired_up = check.generate.fill value(1.0) : tensor<9728xf32> + kernel.launch @ggml_q1_0_q8_1_x4_generic_control[%input_size, %output_size](%input_size, %output_size, %q8_input, %gate_weight, %generic_gate) : [index, index](index, index, tensor<20x144xi8>, tensor<9728x20x18xi8>, tensor<9728xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_generic_control[%input_size, %output_size](%input_size, %output_size, %q8_input, %up_weight, %generic_up) : [index, index](index, index, tensor<20x144xi8>, tensor<9728x20x18xi8>, tensor<9728xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed_pair[%input_size, %output_size](%input_size, %output_size, %q8_input, %gate_weight, %up_weight, %paired_gate, %paired_up) : [index, index](index, index, tensor<20x144xi8>, tensor<9728x20x18xi8>, tensor<9728x20x18xi8>, tensor<9728xf32>, tensor<9728xf32>) + check.expect.close actual(%paired_gate) expected(%generic_gate) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<9728xf32> + check.expect.close actual(%paired_up) expected(%generic_up) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<9728xf32> + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_bonsai_2560x4096_generic_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_q8_1_x4_generic_control[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed_rows2_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed_rows2[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %output) : [index, index](index, index, tensor<20x144xi8>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_shape_differential_case { + %input_size = check.param.choice values([2560, 4096, 9728]) name("input_size") : index + %output_size = check.param.choice values([1024, 2560, 4096, 151669]) name("output_size") : index + %block_count = check.param.choice values([20, 32, 76]) name("block_count") : index + %q8_input = check.generate.fill value(61) : tensor<[%block_count]x144xi8> + %weight = check.generate.fill value(62) : tensor<[%output_size]x[%block_count]x18xi8> + %packed_output = check.generate.fill value(1.0) : tensor<[%output_size]xf32> + %rows2_output = check.generate.fill value(2.0) : tensor<[%output_size]xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %packed_output) : [index, index](index, index, tensor<[%block_count]x144xi8>, tensor<[%output_size]x[%block_count]x18xi8>, tensor<[%output_size]xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed_rows2[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %rows2_output) : [index, index](index, index, tensor<[%block_count]x144xi8>, tensor<[%output_size]x[%block_count]x18xi8>, tensor<[%output_size]xf32>) + check.expect.close actual(%rows2_output) expected(%packed_output) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<[%output_size]xf32> + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_shape_packed_case { + %input_size = check.param.choice values([2560, 4096, 9728]) name("input_size") : index + %output_size = check.param.choice values([1024, 2560, 4096, 151669]) name("output_size") : index + %block_count = check.param.choice values([20, 32, 76]) name("block_count") : index + %q8_input = check.generate.fill value(61) : tensor<[%block_count]x144xi8> + %weight = check.generate.fill value(62) : tensor<[%output_size]x[%block_count]x18xi8> + %output = check.generate.fill value(0.0) : tensor<[%output_size]xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %output) : [index, index](index, index, tensor<[%block_count]x144xi8>, tensor<[%output_size]x[%block_count]x18xi8>, tensor<[%output_size]xf32>) + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_shape_packed_rows2_case { + %input_size = check.param.choice values([2560, 4096, 9728]) name("input_size") : index + %output_size = check.param.choice values([1024, 2560, 4096, 151669]) name("output_size") : index + %block_count = check.param.choice values([20, 32, 76]) name("block_count") : index + %q8_input = check.generate.fill value(61) : tensor<[%block_count]x144xi8> + %weight = check.generate.fill value(62) : tensor<[%output_size]x[%block_count]x18xi8> + %output = check.generate.fill value(0.0) : tensor<[%output_size]xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed_rows2[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %output) : [index, index](index, index, tensor<[%block_count]x144xi8>, tensor<[%output_size]x[%block_count]x18xi8>, tensor<[%output_size]xf32>) + check.return +} + +check.case public @ggml_q1_0_q8_1_x4_vocab_differential_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(151669) : index + %q8_input = check.generate.fill value(61) : tensor<20x144xi8> + %weight = check.generate.fill value(62) : tensor<151669x20x18xi8> + %packed_output = check.generate.fill value(1.0) : tensor<151669xf32> + %rows2_output = check.generate.fill value(2.0) : tensor<151669xf32> + kernel.launch @ggml_q1_0_q8_1_x4_packed[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %packed_output) : [index, index](index, index, tensor<20x144xi8>, tensor<151669x20x18xi8>, tensor<151669xf32>) + kernel.launch @ggml_q1_0_q8_1_x4_packed_rows2[%input_size, %output_size](%input_size, %output_size, %q8_input, %weight, %rows2_output) : [index, index](index, index, tensor<20x144xi8>, tensor<151669x20x18xi8>, tensor<151669xf32>) + check.expect.close actual(%rows2_output) expected(%packed_output) atol(0.050000000000000003) rtol(1.0000000000000001e-05) nan(same) : tensor<151669xf32> + check.return +} + +check.benchmark<@ggml_q1_0_q8_1_x4_bonsai_2560x4096_generic_case> @ggml_q1_0_q8_1_x4_bonsai_2560x4096_generic + +check.benchmark<@ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed_case> @ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed + +check.benchmark<@ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed_rows2_case> @ggml_q1_0_q8_1_x4_bonsai_2560x4096_packed_rows2 + +check.benchmark<@ggml_q1_0_q8_1_x4_bonsai_2560x4096_differential_case> @ggml_q1_0_q8_1_x4_bonsai_2560x4096_correctness + +check.benchmark<@ggml_q1_0_q8_1_x4_bonsai_2560x9728_pair_differential_case> @ggml_q1_0_q8_1_x4_bonsai_2560x9728_pair_correctness + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_differential_case> @ggml_q1_0_q8_1_x4_2560x1024_correctness {block_count = 20, input_size = 2560, output_size = 1024} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_differential_case> @ggml_q1_0_q8_1_x4_4096x2560_correctness {block_count = 32, input_size = 4096, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_differential_case> @ggml_q1_0_q8_1_x4_9728x2560_correctness {block_count = 76, input_size = 9728, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_vocab_differential_case> @ggml_q1_0_q8_1_x4_2560x151669_correctness + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_case> @ggml_q1_0_q8_1_x4_2560x1024_packed {block_count = 20, input_size = 2560, output_size = 1024} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_rows2_case> @ggml_q1_0_q8_1_x4_2560x1024_packed_rows2 {block_count = 20, input_size = 2560, output_size = 1024} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_case> @ggml_q1_0_q8_1_x4_4096x2560_packed {block_count = 32, input_size = 4096, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_rows2_case> @ggml_q1_0_q8_1_x4_4096x2560_packed_rows2 {block_count = 32, input_size = 4096, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_case> @ggml_q1_0_q8_1_x4_9728x2560_packed {block_count = 76, input_size = 9728, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_rows2_case> @ggml_q1_0_q8_1_x4_9728x2560_packed_rows2 {block_count = 76, input_size = 9728, output_size = 2560} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_case> @ggml_q1_0_q8_1_x4_2560x151669_packed {block_count = 20, input_size = 2560, output_size = 151669} + +check.benchmark<@ggml_q1_0_q8_1_x4_shape_packed_rows2_case> @ggml_q1_0_q8_1_x4_2560x151669_packed_rows2 {block_count = 20, input_size = 2560, output_size = 151669} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_vector_schedule_test.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_vector_schedule_test.loom new file mode 100644 index 000000000000..453f7b0660a2 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/q1_0_vector_schedule_test.loom @@ -0,0 +1,476 @@ +// Test-only controls for the Q1_0 vector schedule used by Vulkan. +func.decl @ggml_q1_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) + +amdgpu.target @ggml_q1_0_vector_schedule_gfx11_wave32 {subgroup_size = 32} + +func.def inline @ggml_q1_0_f32_x4_row_lane(%input_size: index, %input: buffer, %weight: buffer, %weight_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%bounded_input_size]xf32> + %iteration_count = index.div %bounded_input_size, %c128 : index + %lane_k = index.mul %bounded_lane, %c4 : index + %last_vector_origin = index.sub %bounded_input_size, %c4 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%acc = %c0_f32 : f32) -> (f32) unroll(%c4) { + %iteration_k = index.mul %iteration, %c128 : index + %k0 = index.add %iteration_k, %lane_k : index + %k = index.assume %k0 [range(%k0, 0, %last_vector_origin)] : index + %weights = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %mask = vector.mask.range [%k to %bounded_input_size step %c1] : index -> vector<4xi1> + %activations = vector.load.mask %input_view[%k], %mask, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %products = vector.mulf %weights, %activations : vector<4xf32> + %contribution = vector.reduce %products, %c0_f32 : vector<4xf32>, f32 + %next = scalar.addf %acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q1_0_f32_x8_row_lane(%input_size: index, %input: buffer, %weight: buffer, %weight_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%bounded_input_size]xf32> + %iteration_count = index.div %bounded_input_size, %c256 : index + %lane_k = index.mul %bounded_lane, %c8 : index + %last_vector_origin = index.sub %bounded_input_size, %c4 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%acc = %c0_f32 : f32) -> (f32) unroll(%c4) { + %iteration_k = index.mul %iteration, %c256 : index + %k00 = index.add %iteration_k, %lane_k : index + %k10 = index.add %k00, %c4 : index + %k0 = index.assume %k00 [range(%k00, 0, %last_vector_origin)] : index + %k1 = index.assume %k10 [range(%k10, 0, %last_vector_origin)] : index + %weights0 = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %weights1 = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %mask0 = vector.mask.range [%k0 to %bounded_input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %bounded_input_size step %c1] : index -> vector<4xi1> + %activations0 = vector.load.mask %input_view[%k0], %mask0, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %activations1 = vector.load.mask %input_view[%k1], %mask1, %c0_f32x4 : view<[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %products0 = vector.mulf %weights0, %activations0 : vector<4xf32> + %products1 = vector.mulf %weights1, %activations1 : vector<4xf32> + %contribution0 = vector.reduce %products0, %c0_f32 : vector<4xf32>, f32 + %contribution1 = vector.reduce %products1, %c0_f32 : vector<4xf32>, f32 + %pair = scalar.addf %contribution0, %contribution1 : f32 + %next = scalar.addf %acc, %pair : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q1_0_f16_x8_row_lane(%input_size: index, %input: buffer, %weight: buffer, %weight_row_byte_base: offset, %lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%bounded_input_size]xf16> + %iteration_count = index.div %bounded_input_size, %c256 : index + %lane_k = index.mul %bounded_lane, %c8 : index + %last_vector_origin = index.sub %bounded_input_size, %c4 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%acc = %c0_f32 : f32) -> (f32) unroll(%c4) { + %iteration_k = index.mul %iteration, %c256 : index + %k00 = index.add %iteration_k, %lane_k : index + %k10 = index.add %k00, %c4 : index + %k0 = index.assume %k00 [range(%k00, 0, %last_vector_origin)] : index + %k1 = index.assume %k10 [range(%k10, 0, %last_vector_origin)] : index + %weights0 = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %weights1 = func.call @ggml_q1_0_f32_vector4(%weight, %weight_row_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %mask0 = vector.mask.range [%k0 to %bounded_input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %bounded_input_size step %c1] : index -> vector<4xi1> + %activations0_f16 = vector.load.mask %input_view[%k0], %mask0, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %activations1_f16 = vector.load.mask %input_view[%k1], %mask1, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %activations0 = vector.extf %activations0_f16 : vector<4xf16> to vector<4xf32> + %activations1 = vector.extf %activations1_f16 : vector<4xf16> to vector<4xf32> + %products0 = vector.mulf %weights0, %activations0 : vector<4xf32> + %products1 = vector.mulf %weights1, %activations1 : vector<4xf32> + %contribution0 = vector.reduce %products0, %c0_f32 : vector<4xf32>, f32 + %contribution1 = vector.reduce %products1, %c0_f32 : vector<4xf32>, f32 + %pair = scalar.addf %contribution0, %contribution1 : f32 + %next = scalar.addf %acc, %pair : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_q1_0_f16_x8_rows2_lane(%input_size: index, %input: buffer, %weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %lane: index) -> (f32, f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%bounded_input_size]xf16> + %iteration_count = index.div %bounded_input_size, %c256 : index + %lane_k = index.mul %bounded_lane, %c8 : index + %last_vector_origin = index.sub %bounded_input_size, %c4 : index + %result0, %result1 = scf.for %iteration = [%c0 to %iteration_count step %c1](%acc0 = %c0_f32 : f32, %acc1 = %c0_f32 : f32) -> (f32, f32) unroll(%c4) { + %iteration_k = index.mul %iteration, %c256 : index + %k00 = index.add %iteration_k, %lane_k : index + %k10 = index.add %k00, %c4 : index + %k0 = index.assume %k00 [range(%k00, 0, %last_vector_origin)] : index + %k1 = index.assume %k10 [range(%k10, 0, %last_vector_origin)] : index + %mask0 = vector.mask.range [%k0 to %bounded_input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %bounded_input_size step %c1] : index -> vector<4xi1> + %activations0_f16 = vector.load.mask %input_view[%k0], %mask0, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %activations1_f16 = vector.load.mask %input_view[%k1], %mask1, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %activations0 = vector.extf %activations0_f16 : vector<4xf16> to vector<4xf32> + %activations1 = vector.extf %activations1_f16 : vector<4xf16> to vector<4xf32> + %weights00 = func.call @ggml_q1_0_f32_vector4(%weight, %row0_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %weights01 = func.call @ggml_q1_0_f32_vector4(%weight, %row0_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %weights10 = func.call @ggml_q1_0_f32_vector4(%weight, %row1_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %weights11 = func.call @ggml_q1_0_f32_vector4(%weight, %row1_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %p00 = vector.mulf %weights00, %activations0 : vector<4xf32> + %p01 = vector.mulf %weights01, %activations1 : vector<4xf32> + %p10 = vector.mulf %weights10, %activations0 : vector<4xf32> + %p11 = vector.mulf %weights11, %activations1 : vector<4xf32> + %s00 = vector.reduce %p00, %c0_f32 : vector<4xf32>, f32 + %s01 = vector.reduce %p01, %c0_f32 : vector<4xf32>, f32 + %s10 = vector.reduce %p10, %c0_f32 : vector<4xf32>, f32 + %s11 = vector.reduce %p11, %c0_f32 : vector<4xf32>, f32 + %sum0 = scalar.addf %s00, %s01 : f32 + %sum1 = scalar.addf %s10, %s11 : f32 + %next0 = scalar.addf %acc0, %sum0 : f32 + %next1 = scalar.addf %acc1, %sum1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + func.return %result0, %result1 : f32, f32 +} + +func.def inline @ggml_q1_0_f16_x8_rows4_lane(%input_size: index, %input: buffer, %weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %row2_byte_base: offset, %row3_byte_base: offset, %lane: index) -> (f32, f32, f32, f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%bounded_input_size]xf16> + %iteration_count = index.div %bounded_input_size, %c256 : index + %lane_k = index.mul %bounded_lane, %c8 : index + %last_vector_origin = index.sub %bounded_input_size, %c4 : index + %result0, %result1, %result2, %result3 = scf.for %iteration = [%c0 to %iteration_count step %c1](%acc0 = %c0_f32 : f32, %acc1 = %c0_f32 : f32, %acc2 = %c0_f32 : f32, %acc3 = %c0_f32 : f32) -> (f32, f32, f32, f32) unroll(%c4) { + %iteration_k = index.mul %iteration, %c256 : index + %k00 = index.add %iteration_k, %lane_k : index + %k10 = index.add %k00, %c4 : index + %k0 = index.assume %k00 [range(%k00, 0, %last_vector_origin)] : index + %k1 = index.assume %k10 [range(%k10, 0, %last_vector_origin)] : index + %mask0 = vector.mask.range [%k0 to %bounded_input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %bounded_input_size step %c1] : index -> vector<4xi1> + %a0_f16 = vector.load.mask %input_view[%k0], %mask0, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %a1_f16 = vector.load.mask %input_view[%k1], %mask1, %c0_f16x4 : view<[%bounded_input_size]xf16>, vector<4xi1>, vector<4xf16> + %a0 = vector.extf %a0_f16 : vector<4xf16> to vector<4xf32> + %a1 = vector.extf %a1_f16 : vector<4xf16> to vector<4xf32> + %w00 = func.call @ggml_q1_0_f32_vector4(%weight, %row0_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %w01 = func.call @ggml_q1_0_f32_vector4(%weight, %row0_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %w10 = func.call @ggml_q1_0_f32_vector4(%weight, %row1_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %w11 = func.call @ggml_q1_0_f32_vector4(%weight, %row1_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %w20 = func.call @ggml_q1_0_f32_vector4(%weight, %row2_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %w21 = func.call @ggml_q1_0_f32_vector4(%weight, %row2_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %w30 = func.call @ggml_q1_0_f32_vector4(%weight, %row3_byte_base, %k0) : (buffer, offset, index) -> (vector<4xf32>) + %w31 = func.call @ggml_q1_0_f32_vector4(%weight, %row3_byte_base, %k1) : (buffer, offset, index) -> (vector<4xf32>) + %p00 = vector.mulf %w00, %a0 : vector<4xf32> + %p01 = vector.mulf %w01, %a1 : vector<4xf32> + %p10 = vector.mulf %w10, %a0 : vector<4xf32> + %p11 = vector.mulf %w11, %a1 : vector<4xf32> + %p20 = vector.mulf %w20, %a0 : vector<4xf32> + %p21 = vector.mulf %w21, %a1 : vector<4xf32> + %p30 = vector.mulf %w30, %a0 : vector<4xf32> + %p31 = vector.mulf %w31, %a1 : vector<4xf32> + %s00 = vector.reduce %p00, %c0_f32 : vector<4xf32>, f32 + %s01 = vector.reduce %p01, %c0_f32 : vector<4xf32>, f32 + %s10 = vector.reduce %p10, %c0_f32 : vector<4xf32>, f32 + %s11 = vector.reduce %p11, %c0_f32 : vector<4xf32>, f32 + %s20 = vector.reduce %p20, %c0_f32 : vector<4xf32>, f32 + %s21 = vector.reduce %p21, %c0_f32 : vector<4xf32>, f32 + %s30 = vector.reduce %p30, %c0_f32 : vector<4xf32>, f32 + %s31 = vector.reduce %p31, %c0_f32 : vector<4xf32>, f32 + %sum0 = scalar.addf %s00, %s01 : f32 + %sum1 = scalar.addf %s10, %s11 : f32 + %sum2 = scalar.addf %s20, %s21 : f32 + %sum3 = scalar.addf %s30, %s31 : f32 + %next0 = scalar.addf %acc0, %sum0 : f32 + %next1 = scalar.addf %acc1, %sum1 : f32 + %next2 = scalar.addf %acc2, %sum2 : f32 + %next3 = scalar.addf %acc3, %sum3 : f32 + scf.yield %next0, %next1, %next2, %next3 : f32, f32, f32, f32 + } + func.return %result0, %result1, %result2, %result3 : f32, f32, f32, f32 +} + +kernel.def target(@ggml_q1_0_vector_schedule_gfx11_wave32) @ggml_q1_0_f32_x4_rows1(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size = index.assume %row0, %output_size [lt(%row0, %output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 18 : offset + %blocks = index.div %bounded_input_size, %c128 : index + %row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %partial = func.call @ggml_q1_0_f32_x4_row_lane(%bounded_input_size, %input, %weight, %row_base, %lane) : (index, buffer, buffer, offset, index) -> (f32) + %sum = kernel.subgroup.reduce %partial : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum, %output_view[%row] : f32, view<[%bounded_output_size]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_vector_schedule_gfx11_wave32) @ggml_q1_0_f32_x8_rows1(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size = index.assume %row0, %output_size [lt(%row0, %output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 18 : offset + %blocks = index.div %bounded_input_size, %c128 : index + %row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %partial = func.call @ggml_q1_0_f32_x8_row_lane(%bounded_input_size, %input, %weight, %row_base, %lane) : (index, buffer, buffer, offset, index) -> (f32) + %sum = kernel.subgroup.reduce %partial : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum, %output_view[%row] : f32, view<[%bounded_output_size]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_vector_schedule_gfx11_wave32) @ggml_q1_0_f16_x8_rows1(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size = index.assume %row0, %output_size [lt(%row0, %output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 18 : offset + %blocks = index.div %bounded_input_size, %c128 : index + %row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %partial = func.call @ggml_q1_0_f16_x8_row_lane(%bounded_input_size, %input, %weight, %row_base, %lane) : (index, buffer, buffer, offset, index) -> (f32) + %sum = kernel.subgroup.reduce %partial : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum, %output_view[%row] : f32, view<[%bounded_output_size]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_vector_schedule_gfx11_wave32) @ggml_q1_0_f16_x8_rows2(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c32 = index.constant 32 : index + %output_size_rounded = index.add %output_size, %c1 : index + %groups = index.div %output_size_rounded, %c2 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %group = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 18 : offset + %row00 = index.mul %group, %c2 : index + %row0, %bounded_output_size = index.assume %row00, %output_size [lt(%row00, %output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %blocks = index.div %bounded_input_size, %c128 : index + %row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %row0_base = index.scale %row0, %row_bytes : index, offset -> offset + %row1_base = index.scale %row1, %row_bytes : index, offset -> offset + %partial0, %partial1 = func.call @ggml_q1_0_f16_x8_rows2_lane(%bounded_input_size, %input, %weight, %row0_base, %row1_base, %lane) : (index, buffer, buffer, offset, offset, index) -> (f32, f32) + %sum0 = kernel.subgroup.reduce %partial0 : f32 + %sum1 = kernel.subgroup.reduce %partial1 : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum0, %output_view[%row0] : f32, view<[%bounded_output_size]xf32> + %has_row1 = index.cmp ult, %row1, %bounded_output_size : index + scf.if %has_row1 { + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + view.store %sum1, %output_view[%bounded_row1] : f32, view<[%bounded_output_size]xf32> + } + } + kernel.return +} + +kernel.def target(@ggml_q1_0_vector_schedule_gfx11_wave32) @ggml_q1_0_f16_x8_rows4(%input_size: index, %output_size: index) { + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %rounded = index.add %output_size, %c3 : index + %groups = index.div %rounded, %c4 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %group = kernel.workgroup.id : 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 + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %block_bytes = index.constant 18 : offset + %row00 = index.mul %group, %c4 : index + %row0, %bounded_output_size = index.assume %row00, %output_size [lt(%row00, %output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row2 = index.add %row0, %c2 : index + %row3 = index.add %row0, %c3 : index + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + %bounded_row2 = index.assume %row2 [lt(%row2, %bounded_output_size)] : index + %bounded_row3 = index.assume %row3 [lt(%row3, %bounded_output_size)] : index + %blocks = index.div %bounded_input_size, %c128 : index + %row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %row0_base = index.scale %row0, %row_bytes : index, offset -> offset + %row1_base = index.scale %bounded_row1, %row_bytes : index, offset -> offset + %row2_base = index.scale %bounded_row2, %row_bytes : index, offset -> offset + %row3_base = index.scale %bounded_row3, %row_bytes : index, offset -> offset + %partial0, %partial1, %partial2, %partial3 = func.call @ggml_q1_0_f16_x8_rows4_lane(%bounded_input_size, %input, %weight, %row0_base, %row1_base, %row2_base, %row3_base, %lane) : (index, buffer, buffer, offset, offset, offset, offset, index) -> (f32, f32, f32, f32) + %sum0 = kernel.subgroup.reduce %partial0 : f32 + %sum1 = kernel.subgroup.reduce %partial1 : f32 + %sum2 = kernel.subgroup.reduce %partial2 : f32 + %sum3 = kernel.subgroup.reduce %partial3 : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_output_size]xf32> + view.store %sum0, %output_view[%row0] : f32, view<[%bounded_output_size]xf32> + view.store %sum1, %output_view[%bounded_row1] : f32, view<[%bounded_output_size]xf32> + view.store %sum2, %output_view[%bounded_row2] : f32, view<[%bounded_output_size]xf32> + view.store %sum3, %output_view[%bounded_row3] : f32, view<[%bounded_output_size]xf32> + } + kernel.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_differential { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input_f32 = check.generate.fill value(0.25) : tensor<2560xf32> + %input_f16 = check.generate.fill value(0.25) : tensor<2560xf16> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %reference = check.generate.fill value(0.0) : tensor<4096xf32> + %f32_x8 = check.generate.fill value(1.0) : tensor<4096xf32> + %f16_x8 = check.generate.fill value(2.0) : tensor<4096xf32> + %f16_rows2 = check.generate.fill value(3.0) : tensor<4096xf32> + %f16_rows4 = check.generate.fill value(4.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f32_x4_rows1[%input_size, %output_size](%input_size, %output_size, %input_f32, %weight, %reference) : [index, index](index, index, tensor<2560xf32>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_f32_x8_rows1[%input_size, %output_size](%input_size, %output_size, %input_f32, %weight, %f32_x8) : [index, index](index, index, tensor<2560xf32>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_f16_x8_rows1[%input_size, %output_size](%input_size, %output_size, %input_f16, %weight, %f16_x8) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_f16_x8_rows2[%input_size, %output_size](%input_size, %output_size, %input_f16, %weight, %f16_rows2) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + kernel.launch @ggml_q1_0_f16_x8_rows4[%input_size, %output_size](%input_size, %output_size, %input_f16, %weight, %f16_rows4) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.expect.close actual(%f32_x8) expected(%reference) atol(0.10000000000000001) rtol(0.0001) nan(same) : tensor<4096xf32> + check.expect.close actual(%f16_x8) expected(%reference) atol(0.10000000000000001) rtol(0.0001) nan(same) : tensor<4096xf32> + check.expect.close actual(%f16_rows2) expected(%reference) atol(0.10000000000000001) rtol(0.0001) nan(same) : tensor<4096xf32> + check.expect.close actual(%f16_rows4) expected(%reference) atol(0.10000000000000001) rtol(0.0001) nan(same) : tensor<4096xf32> + check.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_f32_x4_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input = check.generate.fill value(0.25) : tensor<2560xf32> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f32_x4_rows1[%input_size, %output_size](%input_size, %output_size, %input, %weight, %output) : [index, index](index, index, tensor<2560xf32>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_f32_x8_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input = check.generate.fill value(0.25) : tensor<2560xf32> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f32_x8_rows1[%input_size, %output_size](%input_size, %output_size, %input, %weight, %output) : [index, index](index, index, tensor<2560xf32>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_f16_x8_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input = check.generate.fill value(0.25) : tensor<2560xf16> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f16_x8_rows1[%input_size, %output_size](%input_size, %output_size, %input, %weight, %output) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows2_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input = check.generate.fill value(0.25) : tensor<2560xf16> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f16_x8_rows2[%input_size, %output_size](%input_size, %output_size, %input, %weight, %output) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.case public @ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows4_case { + %input_size = check.literal value(2560) : index + %output_size = check.literal value(4096) : index + %input = check.generate.fill value(0.25) : tensor<2560xf16> + %weight = check.generate.fill value(62) : tensor<4096x20x18xi8> + %output = check.generate.fill value(0.0) : tensor<4096xf32> + kernel.launch @ggml_q1_0_f16_x8_rows4[%input_size, %output_size](%input_size, %output_size, %input, %weight, %output) : [index, index](index, index, tensor<2560xf16>, tensor<4096x20x18xi8>, tensor<4096xf32>) + check.return +} + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_f32_x4_case> @ggml_q1_0_vector_schedule_2560x4096_f32_x4 + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_f32_x8_case> @ggml_q1_0_vector_schedule_2560x4096_f32_x8 + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_f16_x8_case> @ggml_q1_0_vector_schedule_2560x4096_f16_x8 + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows2_case> @ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows2 + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows4_case> @ggml_q1_0_vector_schedule_2560x4096_f16_x8_rows4 + +check.benchmark<@ggml_q1_0_vector_schedule_2560x4096_differential> @ggml_q1_0_vector_schedule_2560x4096_correctness diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom new file mode 100644 index 000000000000..3683ded16b31 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/swiglu_q1_0_q8_1_x4_f32.loom @@ -0,0 +1,41 @@ +func.decl @ggml_q1_0_q8_1_x4_pair_lane(%input_size: index, %gate: buffer, %gate_row_byte_base: offset, %up: buffer, %up_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) + +amdgpu.target @ggml_q1_0_q8_1_x4_swiglu_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.matmul.q1_0_q8_1_x4.input_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.matmul.q1_0_q8_1_x4.output_size : %value: index where [range(%value, 1, 524288)] + +kernel.def target(@ggml_q1_0_q8_1_x4_swiglu_gfx11_wave32) @ggml_mul_mat_swiglu_q1_0_q8_1_x4_f32(%token_count: index) { + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.q1_0_q8_1_x4.input_size : index + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size = index.assume %row0, %output_size [lt(%row0, %output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %weight_block_bytes = index.constant 18 : offset + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_row_byte_base = index.constant 0 : offset + %gate_lane, %up_lane = func.call @ggml_q1_0_q8_1_x4_pair_lane(%bounded_input_size, %gate_weight, %weight_row_byte_base, %up_weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, buffer, offset, index) -> (f32, f32) + %gate = kernel.subgroup.reduce %gate_lane : f32 + %up = kernel.subgroup.reduce %up_lane : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %sigmoid = scalar.logisticf %gate : f32 + %silu = scalar.mulf %gate, %sigmoid : f32 + %value = scalar.mulf %silu, %up : f32 + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + view.store %value, %output_view[%c0, %row] : f32, view<[%bounded_token_count]x[%bounded_output_size]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_codebook_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_codebook_publish_f32.loom new file mode 100644 index 000000000000..2cc9179ff08a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_codebook_publish_f32.loom @@ -0,0 +1,324 @@ +// Tiled GGML matmul for contiguous 2D F32 activations. +// +// This wrapper owns the single-output tiled kernel family. The tiled +// accumulation core is shared; each kernel symbol selects an epilogue variant. +template.decl @ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_postops.epilogue : %value: index where [range(%value, 1, 3)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %publishes_f16_alternate = index.cmp eq, %epilogue, %c4 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + %output_noalias = buffer.assume.noalias %output : buffer + %partition = kernel.workgroup.id : index + %element_bytes = index.constant 4 : offset + %plane_elements = index.mul %token_count, %output_size : index + %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset + %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset + %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %accumulated = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %accumulated, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %accumulated : vector<4xf32> + } + %residual_added = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + %published = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %residual_added) : (index, vector<4xf32>) -> (vector<4xf32>) + scf.if %publishes_f16_alternate { + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c32, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %output) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c16, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %residual_input) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } else { + scf.if %uses_residual { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + } + template.return +} +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @ggml_matmul_tiled_input_f32_publish_f32_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %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) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %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) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %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) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %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) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %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) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %f16_output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %f16_output, %output, %f16_output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + + + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = config.get @ggml.mul_mat_postops.epilogue : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_codebook_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %codebook, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom index 217155372655..71fca22091be 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom @@ -2,7 +2,9 @@ // // This wrapper owns the single-output tiled kernel family. The tiled // accumulation core is shared; each kernel symbol selects an epilogue variant. -template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32_token64.core(%full_token_tiles: i1, %weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) @@ -43,7 +45,6 @@ config.decl @ggml.mul_mat_postops.epilogue : %value: index where [range(%value, config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { - %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index %c0 = index.constant 0 : index %c1 = index.constant 1 : index @@ -53,6 +54,7 @@ template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 %c16 = index.constant 16 : index %c32 = index.constant 32 : index %c0_offset = index.constant 0 : offset + %f32_element_bytes = index.constant 4 : offset %scratch_d_byte_add = index.constant 1024 : offset %scratch_bytes = index.constant 1152 : offset %c0_f32x4 = vector.constant 0.0 : vector<4xf32> @@ -63,18 +65,14 @@ template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 %output_noalias = buffer.assume.noalias %output : buffer - %partition = kernel.workgroup.id : index - %element_bytes = index.constant 4 : offset - %plane_elements = index.mul %token_count, %output_size : index - %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset - %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset - %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index scf.if %publish_word { - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %row_element_offset = index.mul %token0, %output_size : index + %row_byte_offset = index.scale %row_element_offset, %f32_element_bytes : index, offset -> offset + %output_row = buffer.view %output_noalias[%row_byte_offset] : buffer -> view<[%output_size]xf32> %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> %accumulated = scf.if %accumulates_output -> (vector<4xf32>) { - %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %residual = vector.load.mask %output_row[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> %sum = vector.addf %residual, %values : vector<4xf32> scf.yield %sum : vector<4xf32> } else { @@ -136,11 +134,67 @@ kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_m %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index %weight_format = config.get @ggml.mul_mat.weight_format : index %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant false : i1 %c0_f32 = scalar.constant 0.0 : f32 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32_aligned(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %c1 = index.constant 1 : 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 + %token_tiles = index.div %bounded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant true : i1 + %c0_f32 = scalar.constant 0.0 : f32 + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32_token64_aligned(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 64, 2048), mul(%token_count, 64)] : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %token_tiles = index.div %bounded_token_count, %c64 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 64, 2048), mul(%token_count, 64)] : index + %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant true : i1 + %c0_f32 = scalar.constant 0.0 : f32 + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32_token64.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } @@ -166,16 +220,15 @@ kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_m %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index %weight_format = config.get @ggml.mul_mat.weight_format : index %epilogue = index.constant 4 : index + %full_token_tiles = scalar.constant false : i1 %c0_f32 = scalar.constant 0.0 : f32 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %f16_output, %output, %f16_output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %f16_output, %output, %f16_output, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } - - kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32(%token_count: index) { %output_size = config.get @ggml.mul_mat_postops.output_size : index %token_capacity = config.get @ggml.workload.token_capacity : index @@ -198,10 +251,39 @@ kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_m %c0 = index.constant 0 : index %identity = index.constant 23 : index %epilogue = config.get @ggml.mul_mat_postops.epilogue : index + %full_token_tiles = scalar.constant false : i1 %c0_f32 = scalar.constant 0.0 : f32 %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32_aligned(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %c1 = index.constant 1 : 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 + %token_tiles = index.div %bounded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = config.get @ggml.mul_mat_postops.epilogue : index + %full_token_tiles = scalar.constant true : i1 + %c0_f32 = scalar.constant 0.0 : f32 + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (i1, index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom index 7f5843424db5..d50185e7dbbd 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom @@ -1,7 +1,7 @@ // Generic paired tiled GGML matmul for contiguous 2D F32 activations. // This wrapper composes paired tiled accumulation with pair postprocess, // optional epilogue adds, and F32 publication. -template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) @@ -83,14 +83,40 @@ kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wav %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index %op = config.get @ggml.mul_mat_swiglu.op : index %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } - +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %c1 = index.constant 1 : 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 + %token_tiles = index.div %bounded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant true : i1 + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32(%token_count: index) { %output_size = config.get @ggml.mul_mat_swiglu.output_size : index @@ -114,9 +140,37 @@ kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wav %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index %op = config.get @ggml.mul_mat_swiglu.op : index %epilogue = config.get @ggml.mul_mat_swiglu.epilogue : index + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %residual_output, %residual_input, %residual_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %residual_output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32_aligned(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %c1 = index.constant 1 : 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 + %token_tiles = index.div %bounded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 32, 2048), mul(%token_count, 32)] : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = config.get @ggml.mul_mat_swiglu.epilogue : index + %full_token_tiles = scalar.constant true : i1 + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %bias, %residual_output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom index 43754b72ee8d..78f81d9d6992 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom @@ -1,5 +1,5 @@ // Generic paired tiled GGML matmul with F32 and K16-major F16 publication. -template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%full_token_tiles: i1, %op: index, %epilogue: index, %gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) @@ -12,10 +12,15 @@ template.decl @ggml.normalization.publish_f16.k16_vector4(%values: vector<4xf32> amdgpu.target @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_k16_gfx11_wave64 {subgroup_size = 64} config.decl @ggml.mul_mat_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + config.decl @ggml.mul_mat_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + config.decl @ggml.mul_mat_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + config.decl @ggml.mul_mat_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + config.decl @ggml.mul_mat_swiglu.op : %value: index where [range(%value, 0, 8)] + config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_k16_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { @@ -51,9 +56,10 @@ kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_k16_gfx11 %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index %op = config.get @ggml.mul_mat_swiglu.op : index %epilogue = index.constant 0 : index + %full_token_tiles = scalar.constant false : i1 %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index - template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %f16_output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core>(%full_token_tiles, %op, %epilogue, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output, %output, %output, %f16_output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_codebook_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_codebook_f32_f32.loom new file mode 100644 index 000000000000..1bb644bd2c57 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_codebook_f32_f32.loom @@ -0,0 +1,217 @@ +// Vector-shaped GGML matmul reduction for batch-size-1 F32 outputs. +template.decl @ggml.matmul.vector_publish.next_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) +template.decl @ggml.matmul.vector_codebook_activation.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.unary_f32.apply(%arg0: index, %arg1: f32) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%arg0: index, %arg1: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 1, 524288)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.matmul.vector.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @ggml_matmul_vector_f32_f32_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %output_noalias = buffer.assume.noalias %output : buffer + %partition = kernel.workgroup.id : index + %element_bytes = index.constant 4 : offset + %plane_elements = index.mul %token_count, %output_size : index + %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset + %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset + %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %activated = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %published) : (index, vector<4xf32>) -> (vector<4xf32>) + vector.store.mask %activated, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.apply<@ggml.matmul.vector_publish.next_vector4>(%publish_word, %publish_format, %token_count0, %output_size0, %token0, %channel, %activated, %next_output) : (i1, index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} +template.def<@ggml.matmul.vector_activation.publish_pair2> device @ggml_matmul_vector_f32_f32_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_pair { + %published0 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value0 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value0 : f32 + } + %activated0 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published0) : (index, f32) -> (f32) + view.store %activated0, %output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %published1 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %row1] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value1 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value1 : f32 + } + %activated1 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published1) : (index, f32) -> (f32) + view.store %activated1, %output_view[%token, %row1] : f32, view<[%token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_matmul_vector_f32_f32_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_iq1_s_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_codebook_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %codebook, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_iq1_m_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_codebook_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %codebook, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_iq2_s_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_codebook_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %codebook, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_iq3_xxs_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_codebook_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %codebook, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_iq3_s_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %codebook: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_codebook_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %codebook, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q1_0_q8_1_x4_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q1_0_q8_1_x4_f32.loom new file mode 100644 index 000000000000..2df32dce63ba --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q1_0_q8_1_x4_f32.loom @@ -0,0 +1,91 @@ +func.decl @ggml_q1_0_q8_1_x4_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +func.decl @ggml_q1_0_q8_1_x4_rows2_lane(%input_size: index, %weight: buffer, %row0_byte_base: offset, %row1_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32, f32) + +amdgpu.target @ggml_q1_0_q8_1_x4_vector_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.matmul.q1_0_q8_1_x4.input_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.matmul.q1_0_q8_1_x4.output_size : %value: index where [range(%value, 1, 524288)] + +kernel.def target(@ggml_q1_0_q8_1_x4_vector_gfx11_wave32) @ggml_mul_mat_vector_q1_0_q8_1_x4_f32(%token_count: index) { + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%output_size, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.q1_0_q8_1_x4.input_size : index + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %row0 = kernel.workgroup.id : index + %row, %bounded_output_size = index.assume %row0, %output_size [lt(%row0, %output_size)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %weight_block_bytes = index.constant 18 : offset + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_row_byte_base = index.constant 0 : offset + %lane_value = func.call @ggml_q1_0_q8_1_x4_row_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %dot = kernel.subgroup.reduce %lane_value : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + view.store %dot, %output_view[%c0, %row] : f32, view<[%bounded_token_count]x[%bounded_output_size]xf32> + } + kernel.return +} + +kernel.def target(@ggml_q1_0_q8_1_x4_vector_gfx11_wave32) @ggml_mul_mat_vector_q1_0_q8_1_x4_rows2_f32(%token_count: index) { + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c32 = index.constant 32 : index + %rounded = index.add %output_size, %c1 : index + %groups = index.div %rounded, %c2 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.q1_0_q8_1_x4.input_size : index + %output_size = config.get @ggml.matmul.q1_0_q8_1_x4.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 128, 32768), mul(%input_size, 128)] : index + %group = kernel.workgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %weight_block_bytes = index.constant 18 : offset + %q8_row_byte_base = index.constant 0 : offset + %row00 = index.mul %group, %c2 : index + %row0, %bounded_output_size = index.assume %row00, %output_size [lt(%row00, %output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %has_row1 = index.cmp ult, %row1, %bounded_output_size : index + %weight_block_count = index.div %bounded_input_size, %c128 : index + %weight_row_bytes = index.scale %weight_block_count, %weight_block_bytes : index, offset -> offset + %row0_base = index.scale %row0, %weight_row_bytes : index, offset -> offset + %partial0, %partial1 = scf.if %has_row1 -> (f32, f32) { + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + %row1_base = index.scale %bounded_row1, %weight_row_bytes : index, offset -> offset + %pair0, %pair1 = func.call @ggml_q1_0_q8_1_x4_rows2_lane(%bounded_input_size, %weight, %row0_base, %row1_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, offset, buffer, offset, index) -> (f32, f32) + scf.yield %pair0, %pair1 : f32, f32 + } else { + %single = func.call @ggml_q1_0_q8_1_x4_row_lane(%bounded_input_size, %weight, %row0_base, %q8_input, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %single, %c0_f32 : f32, f32 + } + %sum0 = kernel.subgroup.reduce %partial0 : f32 + %sum1 = kernel.subgroup.reduce %partial1 : f32 + %leader = index.cmp eq, %lane, %c0 : index + scf.if %leader { + %output_view = buffer.view %output[%q8_row_byte_base] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + view.store %sum0, %output_view[%c0, %row0] : f32, view<[%bounded_token_count]x[%bounded_output_size]xf32> + scf.if %has_row1 { + %bounded_row1 = index.assume %row1 [lt(%row1, %bounded_output_size)] : index + view.store %sum1, %output_view[%c0, %bounded_row1] : f32, view<[%bounded_token_count]x[%bounded_output_size]xf32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom similarity index 76% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom index d8eda208ccc9..a68764d6b80e 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom @@ -6,13 +6,18 @@ template.decl @ggml.rmsnorm_binary_q8_1_x4.body(%publish_f16: i1, %token_count: amdgpu.target @ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32 {subgroup_size = 32} config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_binary_q8_1_x4.publish_f16 : %value: index where [range(%value, 0, 1)] + kernel.decl @ggml_rmsnorm_binary_f32(%token_count$binary: index) launch(%token_count$binary_launch: index, %input$binary: buffer, %rhs$binary: buffer, %output$binary: buffer) + kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$quantize: index, %input_size$quantize: index) launch(%token_count$quantize_launch: index, %input_size$quantize_launch: index, %input$quantize: buffer, %output$quantize: buffer) -kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnorm_binary_q8_1_x4") @ggml_rmsnorm_binary_q8_1_x4(%token_count: index) { +kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnorm_binary_q8_1_x4_publish") @ggml_rmsnorm_binary_q8_1_x4_publish(%token_count: index) { %c1 = index.constant 1 : index %c256 = index.constant 256 : index %hidden = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index @@ -25,13 +30,15 @@ kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnor %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index %groups = index.mul %token_count, %groups_per_row : index kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 2048)] { - %publish_f16 = scalar.constant false : i1 - template.apply<@ggml.rmsnorm_binary_q8_1_x4.body>(%publish_f16, %token_count, %input, %rhs, %output, %q8_output, %output) : (i1, index, buffer, buffer, buffer, buffer, buffer) +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) where [range(%token_count, 1, 2048)] { + %publish_f16_value = config.get @ggml.rmsnorm_binary_q8_1_x4.publish_f16 : index + %c1 = index.constant 1 : index + %publish_f16 = index.cmp eq, %publish_f16_value, %c1 : index + template.apply<@ggml.rmsnorm_binary_q8_1_x4.body>(%publish_f16, %token_count, %input, %rhs, %output, %q8_output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) kernel.return } -check.case public @ggml_rmsnorm_binary_q8_1_x4_mul_case { +check.case public @ggml_rmsnorm_binary_q8_1_x4_publish_mul_case { %token_count = check.literal value(1) : index %hidden_size = check.literal value(128) : index %input = check.generate.fill value(2.0) : tensor<128xf32> @@ -42,7 +49,7 @@ check.case public @ggml_rmsnorm_binary_q8_1_x4_mul_case { kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %f32) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %f32, %expected) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) %actual_f32 = check.generate.fill value(0.0) : tensor<128xf32> - kernel.launch @ggml_rmsnorm_binary_q8_1_x4[%token_count](%token_count, %input, %rhs, %actual_f32, %actual) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<144xi8>) + kernel.launch @ggml_rmsnorm_binary_q8_1_x4_publish[%token_count](%token_count, %input, %rhs, %actual_f32, %actual, %actual) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<144xi8>, tensor<144xi8>) check.expect.equal actual(%actual_f32) expected(%f32) : tensor<128xf32> check.expect.equal actual(%actual) expected(%expected) : tensor<144xi8> check.return diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom deleted file mode 100644 index ba07f1f7a6eb..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_f16.loom +++ /dev/null @@ -1,29 +0,0 @@ -// Copyright 2026 The HRX Authors -// SPDX-License-Identifier: Apache-2.0 - -template.decl @ggml.rmsnorm_binary_q8_1_x4.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) - -amdgpu.target @ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] -config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 -config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] - -kernel.def target(@ggml_rmsnorm_binary_q8_1_x4_gfx11_wave32) export("ggml_rmsnorm_binary_q8_1_x4_f16") @ggml_rmsnorm_binary_q8_1_x4_f16(%token_count: index) { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %hidden = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index - %tile_width = index.constant 1024 : index - %tile_round = index.constant 1023 : index - %decode_limit = index.constant 6 : index - %rounded_hidden = index.add %hidden, %tile_round : index - %stripe_count = index.div %rounded_hidden, %tile_width : index - %is_decode = index.cmp ult, %token_count, %decode_limit : index - %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index - %groups = index.mul %token_count, %groups_per_row : index - kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) where [range(%token_count, 1, 2048)] { - %publish_f16 = scalar.constant true : i1 - template.apply<@ggml.rmsnorm_binary_q8_1_x4.body>(%publish_f16, %token_count, %input, %rhs, %output, %q8_output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom similarity index 70% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom index 6b2e6bba74d1..64da5ef0bf73 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom @@ -6,11 +6,16 @@ template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] + config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] + config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_q8_1_x4") @ggml_rmsnorm_gate_f32_q8_1_x4(%token_count: index) { +config.decl @ggml.rmsnorm_gate_f32.publish_q8 : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_publish") @ggml_rmsnorm_gate_f32_publish(%token_count: index) { %one = index.constant 1 : index %eight = index.constant 8 : index %seven = index.constant 7 : index @@ -18,8 +23,10 @@ kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f %round = index.add %token_count, %seven : index %groups = index.div %round, %eight : index kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %q8_output: buffer) { - %publish_q8 = scalar.constant true : i1 - template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %q8_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) +} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) { + %publish_format = config.get @ggml.rmsnorm_gate_f32.publish_q8 : index + %q8_format = index.constant 1 : index + %publish_q8 = index.cmp eq, %publish_format, %q8_format : index + template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %alternate_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_f16.loom deleted file mode 100644 index af0592a1fb39..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish_f16.loom +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright 2026 The HRX Authors -// SPDX-License-Identifier: Apache-2.0 - -template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) - -amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] -config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 -config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] -config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] - -kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_f32_f16") @ggml_rmsnorm_gate_f32_f16(%token_count: index) { - %one = index.constant 1 : index - %eight = index.constant 8 : index - %seven = index.constant 7 : index - %threads = index.constant 256 : index - %round = index.add %token_count, %seven : index - %groups = index.div %round, %eight : index - kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index -} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %f16_output: buffer) { - %publish_q8 = scalar.constant false : i1 - template.apply<@ggml.rmsnorm_gate_f32.publish_body>(%publish_q8, %token_count, %input, %weight, %raw_gate, %output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom index f3e3b3b874ed..b321fff4e636 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 // Generic GGML ROPE for F32 head-major rows. -template.decl @ggml.rope_f32.body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) +template.decl @ggml.rope_f32.body(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) amdgpu.target @ggml_rope_f32_gfx11_wave32 {subgroup_size = 32} @@ -16,7 +16,7 @@ config.decl @ggml.rope_f32.input_span : %value: index where [range(%value, 1, 10 config.decl @ggml.rope_f32.mscale : f32 config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] -kernel.def target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count: index) { +kernel.def target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count: index, %input_span: index) { %token_capacity = config.get @ggml.rope_f32.token_capacity : index %head_count = config.get @ggml.rope_f32.head_count : index %head_size = config.get @ggml.rope_f32.head_size : index @@ -24,7 +24,7 @@ kernel.def target(@ggml_rope_f32_gfx11_wave32) @ggml_rope_f32(%token_count: inde %c4 = index.constant 4 : index %pair_packet_count = index.div %head_size, %c4 : index kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index -} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - template.apply<@ggml.rope_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %output) : (index, buffer, buffer, buffer, buffer, buffer) +} launch(%token_count: index, %input_span: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%input_span, 1, 1073741824)] { + template.apply<@ggml.rope_f32.body>(%token_count, %input_span, %positions, %input, %theta, %freq_factors, %output) : (index, index, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom index c075f8841c8f..aab4e830eb2e 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom @@ -1,7 +1,7 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @llm.ssm_conv.generic_apply(%linear: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) +template.decl @llm.ssm_conv.generic_apply(%channel: index, %token: index, %sequence: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) @@ -9,10 +9,6 @@ config.decl @llm.ssm_conv.generic.d_conv : %value: index where [range(%value, 1, config.decl @llm.ssm_conv.generic.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] -config.decl @llm.ssm_conv.generic.n_t : %value: index where [range(%value, 1, 512)] - -config.decl @llm.ssm_conv.generic.n_s : %value: index where [range(%value, 1, 4)] - config.decl @llm.ssm_conv.generic.unary_op : %value: index where [range(%value, 0, 23)] config.decl @llm.ssm_conv.generic.binary_op : %value: index where [range(%value, 0, 8)] @@ -21,51 +17,50 @@ config.decl @llm.ssm_conv.generic.binary_lhs : %value: index where [range(%value config.decl @llm.ssm_conv.generic.workgroup_size : %value: index where [range(%value, 256, 256)] -kernel.def export("llm_ssm_conv_binary_f32") @llm_ssm_conv_binary_f32() { +kernel.def export("llm_ssm_conv_binary_f32") @llm_ssm_conv_binary_f32(%n_t: index, %n_s: index) { %c1 = index.constant 1 : index %cneg1 = index.constant -1 : index %d_inner = config.get @llm.ssm_conv.generic.d_inner : index - %n_t = config.get @llm.ssm_conv.generic.n_t : index - %n_s = config.get @llm.ssm_conv.generic.n_s : index + %bounded_n_t = index.assume %n_t [range(%n_t, 1, 512)] : index + %bounded_n_s = index.assume %n_s [range(%n_s, 1, 4)] : index %wg = config.get @llm.ssm_conv.generic.workgroup_size : index - %tokens = index.mul %d_inner, %n_t : index - %total = index.mul %tokens, %n_s : index %rounding = index.add %wg, %cneg1 : index - %rounded = index.add %total, %rounding : index - %groups = index.div %rounded, %wg : index - kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%wg, %c1, %c1) : index -} launch(%window: buffer, %filter: buffer, %operand: buffer, %output: buffer) { + %rounded = index.add %d_inner, %rounding : index + %channel_groups = index.div %rounded, %wg : index + kernel.launch.config workgroups(%channel_groups, %bounded_n_t, %bounded_n_s) workgroup_size(%wg, %c1, %c1) : index +} launch(%n_t: index, %n_s: index, %window: buffer, %filter: buffer, %operand: buffer, %output: buffer) { %base = index.constant 0 : offset - %c0 = index.constant 0 : index %c1 = index.constant 1 : index - %c0_f32 = scalar.constant 0.0 : f32 %d_conv0 = config.get @llm.ssm_conv.generic.d_conv : index %d_inner0 = config.get @llm.ssm_conv.generic.d_inner : index - %n_t0 = config.get @llm.ssm_conv.generic.n_t : index - %n_s0 = config.get @llm.ssm_conv.generic.n_s : index %unary_op = config.get @llm.ssm_conv.generic.unary_op : index %binary_op = config.get @llm.ssm_conv.generic.binary_op : index %binary_lhs = config.get @llm.ssm_conv.generic.binary_lhs : index %wg = config.get @llm.ssm_conv.generic.workgroup_size : index %d_conv = index.assume %d_conv0 [range(%d_conv0, 1, 16)] : index %d_inner = index.assume %d_inner0 [range(%d_inner0, 32, 65536), mul(%d_inner0, 32)] : index - %n_t = index.assume %n_t0 [range(%n_t0, 1, 512)] : index - %n_s = index.assume %n_s0 [range(%n_s0, 1, 4)] : index - - %token_span0 = index.mul %d_inner, %n_t : index - %token_span = index.assume %token_span0 [range(%token_span0, 32, 33554432), mul(%token_span0, 32)] : index - %total0 = index.mul %token_span, %n_s : index - %total = index.assume %total0 [range(%total0, 32, 134217728), mul(%total0, 32)] : index + %bounded_n_t = index.assume %n_t [range(%n_t, 1, 512)] : index %group0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %sequence0 = kernel.workgroup.id : index %lane0 = kernel.workitem.id : index - %group = index.assume %group0 [range(%group0, 0, 524287)] : index + %group = index.assume %group0 [range(%group0, 0, 255)] : index + %token = index.assume %token0 [range(%token0, 0, 511)] : index + %sequence = index.assume %sequence0 [range(%sequence0, 0, 3)] : index %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index - %base_idx0 = index.mul %group, %wg : index - %linear0 = index.add %base_idx0, %lane : index - %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index - %in_bounds = index.cmp ult, %linear, %total : index + %channel_base = index.mul %group, %wg : index + %channel0 = index.add %channel_base, %lane : index + %channel = index.assume %channel0 [range(%channel0, 0, 65791)] : index + %in_bounds = index.cmp ult, %channel, %d_inner : index + %token_span0 = index.mul %d_inner, %bounded_n_t : index + %token_span = index.assume %token_span0 [range(%token_span0, 32, 33554432), mul(%token_span0, 32)] : index + %sequence_base = index.mul %sequence, %token_span : index + %token_base = index.mul %token, %d_inner : index + %row_base = index.add %sequence_base, %token_base : index + %linear0 = index.add %row_base, %channel : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217727)] : index %window_g = buffer.assume.memory_space %window : buffer %filter_g = buffer.assume.memory_space %filter : buffer @@ -78,7 +73,7 @@ kernel.def export("llm_ssm_conv_binary_f32") @llm_ssm_conv_binary_f32() { %output_v = buffer.view %output_na[%base] : buffer -> view<134217728xf32> scf.if %in_bounds { - %activated = template.apply<@llm.ssm_conv.generic_apply>(%linear, %d_conv, %d_inner, %n_t, %unary_op, %window_v, %filter_v) : (index, index, index, index, index, view<1073741824xf32>, view<1048576xf32>) -> (f32) + %activated = template.apply<@llm.ssm_conv.generic_apply>(%channel, %token, %sequence, %d_conv, %d_inner, %bounded_n_t, %unary_op, %window_v, %filter_v) : (index, index, index, index, index, index, index, view<1073741824xf32>, view<1048576xf32>) -> (f32) %operand_value = view.load %operand_v[%linear] : view<134217728xf32> -> f32 %is_lhs = index.cmp eq, %binary_lhs, %c1 : index %lhs = scf.select %is_lhs, %activated, %operand_value : f32 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_f32.loom index d58553c00ba2..a8ad12a57057 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_f32.loom @@ -1,63 +1,57 @@ // Copyright 2026 The HRX Authors // SPDX-License-Identifier: Apache-2.0 -template.decl @llm.ssm_conv.generic_apply(%linear: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) +template.decl @llm.ssm_conv.generic_apply(%channel: index, %token: index, %sequence: index, %d_conv: index, %d_inner: index, %n_t: index, %unary_op: index, %window: view<1073741824xf32>, %filter: view<1048576xf32>) -> (f32) config.decl @llm.ssm_conv.generic.d_conv : %value: index where [range(%value, 1, 16)] config.decl @llm.ssm_conv.generic.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] -config.decl @llm.ssm_conv.generic.n_t : %value: index where [range(%value, 1, 512)] - -config.decl @llm.ssm_conv.generic.n_s : %value: index where [range(%value, 1, 4)] - config.decl @llm.ssm_conv.generic.unary_op : %value: index where [range(%value, 0, 23)] config.decl @llm.ssm_conv.generic.workgroup_size : %value: index where [range(%value, 256, 256)] -kernel.def export("llm_ssm_conv_f32") @llm_ssm_conv_f32() { +kernel.def export("llm_ssm_conv_f32") @llm_ssm_conv_f32(%n_t: index, %n_s: index) { %c1 = index.constant 1 : index %cneg1 = index.constant -1 : index %d_inner = config.get @llm.ssm_conv.generic.d_inner : index - %n_t = config.get @llm.ssm_conv.generic.n_t : index - %n_s = config.get @llm.ssm_conv.generic.n_s : index + %bounded_n_t = index.assume %n_t [range(%n_t, 1, 512)] : index + %bounded_n_s = index.assume %n_s [range(%n_s, 1, 4)] : index %wg = config.get @llm.ssm_conv.generic.workgroup_size : index - %tokens = index.mul %d_inner, %n_t : index - %total = index.mul %tokens, %n_s : index %rounding = index.add %wg, %cneg1 : index - %rounded = index.add %total, %rounding : index - %groups = index.div %rounded, %wg : index - kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%wg, %c1, %c1) : index -} launch(%window: buffer, %filter: buffer, %output: buffer) { + %rounded = index.add %d_inner, %rounding : index + %channel_groups = index.div %rounded, %wg : index + kernel.launch.config workgroups(%channel_groups, %bounded_n_t, %bounded_n_s) workgroup_size(%wg, %c1, %c1) : index +} launch(%n_t: index, %n_s: index, %window: buffer, %filter: buffer, %output: buffer) { %base = index.constant 0 : offset - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c0_f32 = scalar.constant 0.0 : f32 %d_conv0 = config.get @llm.ssm_conv.generic.d_conv : index %d_inner0 = config.get @llm.ssm_conv.generic.d_inner : index - %n_t0 = config.get @llm.ssm_conv.generic.n_t : index - %n_s0 = config.get @llm.ssm_conv.generic.n_s : index %unary_op = config.get @llm.ssm_conv.generic.unary_op : index %wg = config.get @llm.ssm_conv.generic.workgroup_size : index %d_conv = index.assume %d_conv0 [range(%d_conv0, 1, 16)] : index %d_inner = index.assume %d_inner0 [range(%d_inner0, 32, 65536), mul(%d_inner0, 32)] : index - %n_t = index.assume %n_t0 [range(%n_t0, 1, 512)] : index - %n_s = index.assume %n_s0 [range(%n_s0, 1, 4)] : index - - %token_span0 = index.mul %d_inner, %n_t : index - %token_span = index.assume %token_span0 [range(%token_span0, 32, 33554432), mul(%token_span0, 32)] : index - %total0 = index.mul %token_span, %n_s : index - %total = index.assume %total0 [range(%total0, 32, 134217728), mul(%total0, 32)] : index + %bounded_n_t = index.assume %n_t [range(%n_t, 1, 512)] : index %group0 = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %sequence0 = kernel.workgroup.id : index %lane0 = kernel.workitem.id : index - %group = index.assume %group0 [range(%group0, 0, 524287)] : index + %group = index.assume %group0 [range(%group0, 0, 255)] : index + %token = index.assume %token0 [range(%token0, 0, 511)] : index + %sequence = index.assume %sequence0 [range(%sequence0, 0, 3)] : index %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index - %base_idx0 = index.mul %group, %wg : index - %linear0 = index.add %base_idx0, %lane : index - %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index - %in_bounds = index.cmp ult, %linear, %total : index + %channel_base = index.mul %group, %wg : index + %channel0 = index.add %channel_base, %lane : index + %channel = index.assume %channel0 [range(%channel0, 0, 65791)] : index + %in_bounds = index.cmp ult, %channel, %d_inner : index + %token_span0 = index.mul %d_inner, %bounded_n_t : index + %token_span = index.assume %token_span0 [range(%token_span0, 32, 33554432), mul(%token_span0, 32)] : index + %sequence_base = index.mul %sequence, %token_span : index + %token_base = index.mul %token, %d_inner : index + %row_base = index.add %sequence_base, %token_base : index + %linear0 = index.add %row_base, %channel : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217727)] : index %window_g = buffer.assume.memory_space %window : buffer %filter_g = buffer.assume.memory_space %filter : buffer @@ -68,7 +62,7 @@ kernel.def export("llm_ssm_conv_f32") @llm_ssm_conv_f32() { %output_v = buffer.view %output_na[%base] : buffer -> view<134217728xf32> scf.if %in_bounds { - %activated = template.apply<@llm.ssm_conv.generic_apply>(%linear, %d_conv, %d_inner, %n_t, %unary_op, %window_v, %filter_v) : (index, index, index, index, index, view<1073741824xf32>, view<1048576xf32>) -> (f32) + %activated = template.apply<@llm.ssm_conv.generic_apply>(%channel, %token, %sequence, %d_conv, %d_inner, %bounded_n_t, %unary_op, %window_v, %filter_v) : (index, index, index, index, index, index, index, view<1073741824xf32>, view<1048576xf32>) -> (f32) view.store %activated, %output_v[%linear] : f32, view<134217728xf32> } kernel.return diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json index 224ed9a5c8b6..8780c0e17a1e 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/manifest.json @@ -38,6 +38,9 @@ { "path": "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom" }, + { + "path": "../loom-libs/motifs/quantize_q8_1_x4.loom" + }, { "path": "../loom-libs/motifs/quantize_q8_1_x4_group.loom" }, @@ -704,11 +707,70 @@ "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "library_sources": [ - "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ] + }, + "compile_dependencies": [ + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ] + }, + { + "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8", + "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8", + "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", + "primary_sources": [ + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" + ], + "library_sources": [ + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -1022,7 +1084,10 @@ "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "libraries": [ - "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -1264,6 +1329,22 @@ ], "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" }, + { + "name": "routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_plan_test", + "args": [ + "$(location :routed_down_weighted_reduce_next_rmsnorm_f32_linked)", + "--benchmark=@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_differential", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=qwen3_moe.routed_down.output_size=2048", + "--config=qwen3_moe.routed_down.route_count=8", + "--config=qwen3_moe.workload.token_capacity=2", + "--config=ggml.quantize_q8_1_x4.group_capacity=32", + "--dry-run", + "--output-format=jsonl" + ], + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" + }, { "name": "expert_table_partition_fused_plan_test", "args": [ diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom index 0c2850abc6e8..11b35a0ec41c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom @@ -7,6 +7,8 @@ // by the next layer. template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer) +template.decl @ggml.routing.weighted_reduce_next_rmsnorm_f32_publish_q8(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer, %q8_output: buffer) + amdgpu.target @qwen3_moe_routed_down_next_norm_gfx11_wave32 {subgroup_size = 32} config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] @@ -21,6 +23,8 @@ kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count$4: index kernel.decl @qwen3_moe_rmsnorm_f32(%token_count$9: index) launch(%token_count$10: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$11: index, %input_size$12: index) launch(%token_count$13: index, %input_size$14: index, %input: buffer, %output: buffer) + kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32(%token_count: index) { %c1 = index.constant 1 : index %c256 = index.constant 256 : index @@ -30,6 +34,15 @@ kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_rout kernel.return } +kernel.def target(@qwen3_moe_routed_down_next_norm_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer, %next_norm_weight: buffer, %next_projection_input: buffer, %q8_output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.routing.weighted_reduce_next_rmsnorm_f32_publish_q8>(%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input, %q8_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + // The exact Prefill-512 shape compares both published tensors against the // ordinary weighted reduction followed by the canonical RMSNorm. check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential_case { @@ -49,6 +62,53 @@ check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differ check.return } +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_differential_case { + %token_count = check.literal value(2) : index + %hidden_size = check.literal value(2048) : index + %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<2x8xf32> + %routed_output = check.generate.iota offset(-0.25) step(0.015625) period(31) : tensor<2x8x2048xf16> + %expected_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<2x2048xf32> + %actual_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<2x2048xf32> + %next_norm_weight = check.generate.iota offset(0.5) step(0.0078125) period(29) : tensor<2048xf32> + %expected_next_projection_input = check.generate.fill value(0.0) : tensor<2x2048xf32> + %actual_next_projection_input = check.generate.fill value(1.0) : tensor<2x2048xf32> + %expected_q8 = check.generate.fill value(0) : tensor<2x2304xi8> + %actual_q8 = check.generate.fill value(1) : tensor<2x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state, %next_norm_weight, %expected_next_projection_input) : [index](index, tensor<2x8xf32>, tensor<2x8x2048xf16>, tensor<2x2048xf32>, tensor<2048xf32>, tensor<2x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %expected_next_projection_input, %expected_q8) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_next_projection_input, %actual_q8) : [index](index, tensor<2x8xf32>, tensor<2x8x2048xf16>, tensor<2x2048xf32>, tensor<2048xf32>, tensor<2x2048xf32>, tensor<2x2304xi8>) + check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<2x2048xf32> + check.expect.close actual(%actual_next_projection_input) expected(%expected_next_projection_input) atol(0.001) rtol(0.001) nan(same) : tensor<2x2048xf32> + check.expect.equal actual(%actual_q8) expected(%expected_q8) : tensor<2x2304xi8> + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_standalone_q8_benchmark_case { + %token_count = check.literal value(2) : index + %hidden_size = check.literal value(2048) : index + %route_weights = check.generate.fill value(0.125) : tensor<2x8xf32> + %routed_output = check.generate.fill value(0.03125) : tensor<2x8x2048xf16> + %hidden_state = check.generate.fill value(0.25) : tensor<2x2048xf32> + %next_norm_weight = check.generate.fill value(0.5) : tensor<2048xf32> + %next_projection_input = check.generate.fill value(0.0) : tensor<2x2048xf32> + %q8 = check.generate.fill value(0) : tensor<2x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<2x8xf32>, tensor<2x8x2048xf16>, tensor<2x2048xf32>, tensor<2048xf32>, tensor<2x2048xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %next_projection_input, %q8) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) + check.return +} + +check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_benchmark_case { + %token_count = check.literal value(2) : index + %route_weights = check.generate.fill value(0.125) : tensor<2x8xf32> + %routed_output = check.generate.fill value(0.03125) : tensor<2x8x2048xf16> + %hidden_state = check.generate.fill value(0.25) : tensor<2x2048xf32> + %next_norm_weight = check.generate.fill value(0.5) : tensor<2048xf32> + %next_projection_input = check.generate.fill value(0.0) : tensor<2x2048xf32> + %q8 = check.generate.fill value(0) : tensor<2x2304xi8> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input, %q8) : [index](index, tensor<2x8xf32>, tensor<2x8x2048xf16>, tensor<2x2048xf32>, tensor<2048xf32>, tensor<2x2048xf32>, tensor<2x2304xi8>) + check.return +} + check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_benchmark_case { %token_count = check.literal value(512) : index %route_weights = check.generate.fill value(0.125) : tensor<512x8xf32> @@ -77,3 +137,9 @@ check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_be check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_composed_prefill_512 check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differential + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_differential_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_differential + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_standalone_q8_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_standalone_q8_token2 + +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_token2 diff --git a/ggml/src/ggml-hrx/loom-jit.cpp b/ggml/src/ggml-hrx/loom-jit.cpp index 7db1ad29fe33..0afb97e55371 100644 --- a/ggml/src/ggml-hrx/loom-jit.cpp +++ b/ggml/src/ggml-hrx/loom-jit.cpp @@ -13,11 +13,24 @@ #include #include #include +#include #include #include #include #include +struct ggml_hrx_loom_jit_launch_program { + loomc_launch_config_program_t * program = nullptr; + loomc_launch_config_function_t function = loomc_launch_config_function_invalid(); + std::mutex mutex; + + ~ggml_hrx_loom_jit_launch_program() { + if (program != nullptr) { + loomc_launch_config_program_release(program); + } + } +}; + ggml_hrx_loom_jit_compile_result::~ggml_hrx_loom_jit_compile_result() { reset(); } @@ -40,6 +53,7 @@ ggml_hrx_loom_jit_compile_result & ggml_hrx_loom_jit_compile_result::operator=( compile_report_json_size = std::exchange(other.compile_report_json_size, 0); final_module_text = std::exchange(other.final_module_text, nullptr); final_module_text_size = std::exchange(other.final_module_text_size, 0); + launch_program = std::exchange(other.launch_program, nullptr); launch_config = std::exchange(other.launch_config, {}); return *this; } @@ -50,6 +64,7 @@ void ggml_hrx_loom_jit_compile_result::reset() { hrx_host_allocator_free(allocator, manifest_json); hrx_host_allocator_free(allocator, compile_report_json); hrx_host_allocator_free(allocator, final_module_text); + ggml_hrx_loom_jit_launch_program_release(launch_program); hsaco_data = nullptr; hsaco_size = 0; manifest_json = nullptr; @@ -58,6 +73,7 @@ void ggml_hrx_loom_jit_compile_result::reset() { compile_report_json_size = 0; final_module_text = nullptr; final_module_text_size = 0; + launch_program = nullptr; launch_config = {}; } @@ -218,7 +234,6 @@ using LoomResult = LoomHandle using LoomLinkIndexBuilder = LoomHandle; using LoomLinkIndex = LoomHandle; using LoomLinker = LoomHandle; -using LoomLaunchConfigProgram = LoomHandle; struct HrxLoomJitDeleter { void operator()(ggml_hrx_loom_jit_amdgpu * jit) const { ggml_hrx_loom_jit_amdgpu_release(jit); } @@ -435,20 +450,22 @@ hrx_status_t ggml_hrx_loom_jit_copy_artifact_bytes(const loomc_artifact_t * arti return hrx_ok_status(); } -hrx_status_t ggml_hrx_loom_jit_evaluate_launch_config(const loomc_artifact_t * artifact, - const char * root_symbol, - const int64_t * workload_arguments, - size_t workload_argument_count, - ggml_hrx_loom_jit_launch_config * out_launch_config) { - if (!out_launch_config) { - return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_launch_config is required"); +hrx_status_t ggml_hrx_loom_jit_load_launch_program(const loomc_artifact_t * artifact, + const char * root_symbol, + ggml_hrx_loom_jit_launch_program ** out_program) { + if (!out_program) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_program is required"); } + *out_program = nullptr; if (!artifact) { return ggml_hrx_loom_jit_make_status(HRX_STATUS_NOT_FOUND, "Loom did not return a launch-config artifact"); } - LoomLaunchConfigProgram program; - loomc_status_t status = loomc_launch_config_program_load(artifact, loomc_allocator_system(), program.out()); + std::unique_ptr program(new (std::nothrow) ggml_hrx_loom_jit_launch_program()); + if (!program) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_OUT_OF_MEMORY, "failed to allocate Loom launch program"); + } + loomc_status_t status = loomc_launch_config_program_load(artifact, loomc_allocator_system(), &program->program); if (!loomc_status_is_ok(status)) { return ggml_hrx_loom_jit_status_from_loom(status, "load Loom launch config program"); } @@ -457,12 +474,28 @@ hrx_status_t ggml_hrx_loom_jit_evaluate_launch_config(const loomc_artifact_t * if (!export_name.empty() && export_name.front() == '@') { export_name.erase(export_name.begin()); } - loomc_launch_config_function_t function = loomc_launch_config_function_invalid(); - status = loomc_launch_config_program_lookup_function(program.get(), loomc_make_cstring_view(export_name.c_str()), - &function); + status = loomc_launch_config_program_lookup_function(program->program, loomc_make_cstring_view(export_name.c_str()), + &program->function); if (!loomc_status_is_ok(status)) { return ggml_hrx_loom_jit_status_from_loom(status, "find Loom launch config function"); } + *out_program = program.release(); + return hrx_ok_status(); +} + +hrx_status_t ggml_hrx_loom_jit_evaluate_launch_program_impl(ggml_hrx_loom_jit_launch_program * program, + const int64_t * workload_arguments, + size_t workload_argument_count, + ggml_hrx_loom_jit_launch_config * out_launch_config) { + if (!program || !program->program) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "valid launch program is required"); + } + if (!out_launch_config) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "out_launch_config is required"); + } + if (workload_argument_count > 0 && !workload_arguments) { + return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, "workload arguments are required"); + } std::vector workload_bits; workload_bits.reserve(workload_argument_count); @@ -474,9 +507,10 @@ hrx_status_t ggml_hrx_loom_jit_evaluate_launch_config(const loomc_artifact_t * launch_config.type = LOOMC_STRUCTURE_TYPE_LAUNCH_CONFIG; launch_config.structure_size = sizeof(launch_config); - status = loomc_launch_config_program_invoke(program.get(), function, - workload_bits.empty() ? nullptr : workload_bits.data(), - workload_bits.size(), &launch_config); + std::lock_guard lock(program->mutex); + loomc_status_t status = loomc_launch_config_program_invoke( + program->program, program->function, workload_bits.empty() ? nullptr : workload_bits.data(), + workload_bits.size(), &launch_config); if (!loomc_status_is_ok(status)) { return ggml_hrx_loom_jit_status_from_loom(status, "invoke Loom launch config function"); } @@ -551,6 +585,18 @@ hrx_status_t ggml_hrx_loom_jit_parse_sanitizer_reporting(const char * } // namespace +void ggml_hrx_loom_jit_launch_program_release(ggml_hrx_loom_jit_launch_program * program) { + delete program; +} + +hrx_status_t ggml_hrx_loom_jit_launch_program_evaluate(ggml_hrx_loom_jit_launch_program * program, + const int64_t * workload_arguments, + size_t workload_argument_count, + ggml_hrx_loom_jit_launch_config * out_launch_config) { + return ggml_hrx_loom_jit_evaluate_launch_program_impl(program, workload_arguments, workload_argument_count, + out_launch_config); +} + hrx_status_t ggml_hrx_loom_jit_amdgpu_create(const ggml_hrx_loom_jit_amdgpu_options * options, ggml_hrx_loom_jit_amdgpu ** out_jit) { if (!out_jit) { @@ -699,8 +745,8 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * LoomModule module; LoomResult result; std::string specialized_source; - if (options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT && options->dependency_count == 0 && - options->workload_argument_count > 0) { + if (options->specialize_workload && options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_TEXT && + options->dependency_count == 0 && options->workload_argument_count > 0) { std::string specialization_error; const std::string source_text(static_cast(options->source_data), options->source_size); if (!ggml_hrx_loom_specialize_workload_text(source_text, options->root_symbol, options->workload_arguments, @@ -812,9 +858,10 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * LoomSource archived_source; LoomSource specialized_archive_source; std::string specialized_archive_text; - const bool needs_archive_source = + const bool needs_archive_source = options->dependency_count > 0 || - (options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE && options->workload_argument_count > 0); + (options->specialize_workload && options->source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE && + options->workload_argument_count > 0); if (needs_archive_source) { LoomModule archive_module; loomc_link_options_t archive_options = {}; @@ -842,7 +889,7 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * return ggml_hrx_loom_jit_status_from_loom(status, "serialize Loom kernel archive"); } loomc_source_t * archive_index_source = archived_source.get(); - if (options->workload_argument_count > 0) { + if (options->specialize_workload && options->workload_argument_count > 0) { const loomc_byte_span_t archive_contents = loomc_source_contents(archived_source.get()); const std::string archive_text(reinterpret_cast(archive_contents.data), archive_contents.data_length); @@ -928,7 +975,7 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * compile_options.next = &compile_target_options; compile_options.module_name = loomc_make_cstring_view(options->module_name); compile_options.artifact_flags = LOOMC_COMPILE_ARTIFACT_FLAG_MODULE_TEXT | LOOMC_COMPILE_ARTIFACT_FLAG_REPORT_JSON; - if (options->evaluate_launch_config) { + if (options->load_launch_config || options->evaluate_launch_config) { compile_options.artifact_flags |= LOOMC_COMPILE_ARTIFACT_FLAG_LAUNCH_CONFIG; } compile_options.config_flags = LOOMC_CONFIG_POLICY_FLAG_REQUIRE_RESOLVED; @@ -960,7 +1007,7 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * if (!hrx_status_is_ok(hrx_status)) { return hrx_status; } - if (options->evaluate_launch_config) { + if (options->load_launch_config || options->evaluate_launch_config) { const char * launch_config_symbol = options->launch_config_symbol; if (launch_config_symbol == nullptr || launch_config_symbol[0] == '\0') { launch_config_symbol = options->root_symbol; @@ -968,13 +1015,20 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * const loomc_artifact_t * launch_config = ggml_hrx_loom_jit_find_artifact(result.get(), GGML_HRX_LOOM_ARTIFACT_LAUNCH_CONFIG, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_BYTECODE)); - hrx_status = ggml_hrx_loom_jit_evaluate_launch_config( - launch_config, launch_config_symbol, - options->workload_argument_count == 0 ? nullptr : options->workload_arguments, - options->workload_argument_count, &out_result->launch_config); + hrx_status = ggml_hrx_loom_jit_load_launch_program(launch_config, launch_config_symbol, + &out_result->launch_program); if (!hrx_status_is_ok(hrx_status)) { return hrx_status; } + if (options->evaluate_launch_config) { + hrx_status = ggml_hrx_loom_jit_evaluate_launch_program_impl( + out_result->launch_program, + options->workload_argument_count == 0 ? nullptr : options->workload_arguments, + options->workload_argument_count, &out_result->launch_config); + if (!hrx_status_is_ok(hrx_status)) { + return hrx_status; + } + } } result.reset(); diff --git a/ggml/src/ggml-hrx/loom-jit.h b/ggml/src/ggml-hrx/loom-jit.h index 4dd05f73519c..38bf69c43879 100644 --- a/ggml/src/ggml-hrx/loom-jit.h +++ b/ggml/src/ggml-hrx/loom-jit.h @@ -10,6 +10,7 @@ #include struct ggml_hrx_loom_jit_amdgpu; +struct ggml_hrx_loom_jit_launch_program; enum class ggml_hrx_loom_jit_source_format { Text, @@ -61,6 +62,8 @@ struct ggml_hrx_loom_jit_compile_options { size_t config_binding_count = 0; const int64_t * workload_arguments = nullptr; size_t workload_argument_count = 0; + bool specialize_workload = true; + bool load_launch_config = false; bool evaluate_launch_config = false; }; @@ -82,6 +85,7 @@ struct ggml_hrx_loom_jit_compile_result { size_t compile_report_json_size = 0; char * final_module_text = nullptr; size_t final_module_text_size = 0; + ggml_hrx_loom_jit_launch_program * launch_program = nullptr; ggml_hrx_loom_jit_launch_config launch_config; }; @@ -93,3 +97,10 @@ void ggml_hrx_loom_jit_amdgpu_release(ggml_hrx_loom_jit_amdgpu * jit); hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * jit, const ggml_hrx_loom_jit_compile_options * options, ggml_hrx_loom_jit_compile_result * out_result); + +void ggml_hrx_loom_jit_launch_program_release(ggml_hrx_loom_jit_launch_program * program); + +hrx_status_t ggml_hrx_loom_jit_launch_program_evaluate(ggml_hrx_loom_jit_launch_program * program, + const int64_t * workload_arguments, + size_t workload_argument_count, + ggml_hrx_loom_jit_launch_config * out_launch_config); diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 26045fba9cd0..8c30489611a2 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -1219,8 +1219,8 @@ static Status prepare_kernel_command(const CommandProgramExecutionContext & cont } prepared = make_prepared_command_shape(command); - executable_ref = context.kernel_executables->get_or_compile( - { context.device, context.target }, *resolved.definition, dispatch, prepared.kernel.constants); + executable_ref = + context.kernel_executables->get_or_compile({ context.device, context.target }, *resolved.definition, dispatch); if (!executable_ref.valid()) { status.log("failed to prepare %s", command_context.c_str()); return status; @@ -1248,11 +1248,11 @@ static bool execute_prepared_kernel_command(const CommandProgramExecutionContext const KernelExecutable & executable = *command.kernel.executable; hrx_dispatch_config_t config = { - { executable.launch.workgroup_count[0], executable.launch.workgroup_count[1], - executable.launch.workgroup_count[2] }, - { executable.launch.workgroup_size[0], executable.launch.workgroup_size[1], - executable.launch.workgroup_size[2] }, - executable.launch.subgroup_size, + { command.kernel.launch.workgroup_count[0], command.kernel.launch.workgroup_count[1], + command.kernel.launch.workgroup_count[2] }, + { command.kernel.launch.workgroup_size[0], command.kernel.launch.workgroup_size[1], + command.kernel.launch.workgroup_size[2] }, + command.kernel.launch.subgroup_size, }; if (ErrorResult error = take_status(hrx_stream_dispatch( context.stream, executable.executable, executable.export_ordinal, &config, command.kernel.constants.data(), @@ -1290,7 +1290,8 @@ static void materialize_command_list_executables(const CommandProgramExecutionCo for (size_t i = 0; i < prepared_commands.size(); ++i) { PreparedCommand & command = prepared_commands[i]; command.kernel.executable = context.kernel_executables->materialize( - { context.device, context.target }, executable_refs[i], command.kernel.constants); + { context.device, context.target }, executable_refs[i], command.kernel.specialization, + command.kernel.constants, command.kernel.launch); if (command.kernel.executable == nullptr) { status.log("failed to prepare %s", format_prepared_command_context(command).c_str()); } @@ -1582,11 +1583,11 @@ static Status record_prepared_kernel_command(hrx_graph_t graph, executable.executable, executable.export_ordinal, { - { executable.launch.workgroup_count[0], executable.launch.workgroup_count[1], - executable.launch.workgroup_count[2] }, - { executable.launch.workgroup_size[0], executable.launch.workgroup_size[1], - executable.launch.workgroup_size[2] }, - executable.launch.subgroup_size, + { command.kernel.launch.workgroup_count[0], command.kernel.launch.workgroup_count[1], + command.kernel.launch.workgroup_count[2] }, + { command.kernel.launch.workgroup_size[0], command.kernel.launch.workgroup_size[1], + command.kernel.launch.workgroup_size[2] }, + command.kernel.launch.subgroup_size, }, command.kernel.constants.data(), command.kernel.constants.size(), diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index f84ce9fff0b0..bc1debfe03a3 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -4,6 +4,7 @@ #include "dispatch/command-program-resolver.h" #include "dispatch/command-program.h" #include "kernel-corpus/kernel-corpus.h" +#include "loom-jit.h" #include "runtime/graph-replay.h" #include "runtime/host-memory.h" @@ -71,6 +72,7 @@ struct PreparedCommandBinding { struct PreparedKernelCommand { KernelSpecialization specialization; std::shared_ptr executable; + ggml_hrx_loom_jit_launch_config launch; std::vector constants; std::vector bindings; }; diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index ecadaecd9f4f..98d0cff342b6 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -267,7 +267,8 @@ static std::string command_program_shape_key(const CommandProgram & commands) { out << "hrx-command-program-v1|commands=" << commands.commands.size(); for (const Command & command : commands.commands) { out << "|ordinal=" << command.ordinal << "|kind=" << static_cast(command.kind) - << "|kernel=" << command.kernel.kernel_id; + << "|kernel=" << command.kernel.kernel_id + << "|workload=" << static_cast(command.kernel.workload_specialization); for (const auto & parameter : command.kernel.integer_parameters) { out << "|ip:" << parameter.first << '=' << parameter.second; } diff --git a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp index a51f1c4abe81..5a86bc35f616 100644 --- a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.cpp @@ -30,22 +30,69 @@ static void append_u32(std::vector & bytes, uint32_t value) { std::memcpy(bytes.data() + offset, &value, sizeof(value)); } +static void append_u64(std::vector & bytes, uint64_t value) { + const size_t offset = bytes.size(); + bytes.resize(offset + sizeof(value)); + std::memcpy(bytes.data() + offset, &value, sizeof(value)); +} + static bool pack_kernel_constants(const KernelDefinition & definition, - const Dispatch & dispatch, - std::vector & constants) { + const KernelSpecialization & specialization, + size_t expected_size, + std::vector & constants) { constants.clear(); + const size_t parameter_count = definition.launch_parameters.size(); + if (parameter_count == 0) { + return expected_size == 0; + } + if (expected_size % parameter_count != 0) { + GGML_LOG_ERROR("%s: launch scalar ABI size does not match manifest for %s\n", __func__, + kernel_definition_name(definition).c_str()); + return false; + } + const size_t scalar_size = expected_size / parameter_count; + if (scalar_size != sizeof(uint32_t) && scalar_size != sizeof(uint64_t)) { + GGML_LOG_ERROR("%s: unsupported launch scalar ABI width %zu for %s\n", __func__, scalar_size, + kernel_definition_name(definition).c_str()); + return false; + } for (const KernelScalarDefinition & parameter : definition.launch_parameters) { const char * name = parameter.name != nullptr ? parameter.name : ""; const char * type = parameter.type != nullptr ? parameter.type : ""; - const auto item = dispatch.kernel.integer_parameters.find(name); - if (item == dispatch.kernel.integer_parameters.end() || std::strcmp(type, "index") != 0 || item->second < 0 || - static_cast(item->second) > std::numeric_limits::max()) { + const auto item = specialization.integer_parameters.find(name); + if (item == specialization.integer_parameters.end() || std::strcmp(type, "index") != 0 || item->second < 0 || + (scalar_size == sizeof(uint32_t) && + static_cast(item->second) > std::numeric_limits::max())) { constants.clear(); GGML_LOG_ERROR("%s: invalid launch scalar %s for %s\n", __func__, name, kernel_definition_name(definition).c_str()); return false; } - append_u32(constants, static_cast(item->second)); + if (scalar_size == sizeof(uint32_t)) { + append_u32(constants, static_cast(item->second)); + } else { + append_u64(constants, static_cast(item->second)); + } + } + return true; +} + +static bool build_workload_arguments(const KernelDefinition & definition, + const KernelSpecialization & specialization, + std::vector & workload) { + workload.clear(); + workload.reserve(definition.workload_parameters.size()); + for (const KernelScalarDefinition & parameter : definition.workload_parameters) { + const char * name = parameter.name != nullptr ? parameter.name : ""; + const char * type = parameter.type != nullptr ? parameter.type : ""; + const auto item = specialization.integer_parameters.find(name); + if (item == specialization.integer_parameters.end() || std::strcmp(type, "index") != 0) { + GGML_LOG_ERROR("%s: invalid workload scalar %s for %s\n", __func__, name, + kernel_definition_name(definition).c_str()); + workload.clear(); + return false; + } + workload.push_back(item->second); } return true; } @@ -55,15 +102,18 @@ static std::string kernel_executable_key(const KernelDefinition & definition, const char * target) { std::ostringstream out; out << (target != nullptr ? target : "") << '|' << definition.source_digest << '|' << definition.symbol - << "|recipe=" << definition.compile_recipe.mode; - for (const KernelScalarDefinition & parameter : definition.workload_parameters) { - const char * name = parameter.name != nullptr ? parameter.name : ""; - const auto item = dispatch.kernel.integer_parameters.find(name); - out << '|' << name << '='; - if (item == dispatch.kernel.integer_parameters.end()) { - out << ""; - } else { - out << item->second; + << "|recipe=" << definition.compile_recipe.mode << "|workload=" + << (dispatch.kernel.workload_specialization == WorkloadSpecialization::Dynamic ? "dynamic" : "exact"); + if (dispatch.kernel.workload_specialization == WorkloadSpecialization::Exact) { + for (const KernelScalarDefinition & parameter : definition.workload_parameters) { + const char * name = parameter.name != nullptr ? parameter.name : ""; + const auto item = dispatch.kernel.integer_parameters.find(name); + out << '|' << name << '='; + if (item == dispatch.kernel.integer_parameters.end()) { + out << ""; + } else { + out << item->second; + } } } for (const KernelCompileConfig & config : definition.compile_config) { @@ -123,17 +173,9 @@ static bool build_compile_request(const KernelDefinition & definition, request.config_storage.push_back(config); } - request.workload.reserve(definition.workload_parameters.size()); - for (const KernelScalarDefinition & parameter : definition.workload_parameters) { - const char * name = parameter.name != nullptr ? parameter.name : ""; - const char * type = parameter.type != nullptr ? parameter.type : ""; - const auto item = dispatch.kernel.integer_parameters.find(name); - if (item == dispatch.kernel.integer_parameters.end() || std::strcmp(type, "index") != 0) { - GGML_LOG_ERROR("%s: invalid workload scalar %s for %s\n", __func__, name, - kernel_definition_name(definition).c_str()); - return false; - } - request.workload.push_back(item->second); + request.specialize_workload = dispatch.kernel.workload_specialization == WorkloadSpecialization::Exact; + if (!build_workload_arguments(definition, dispatch.kernel, request.workload)) { + return false; } return true; } @@ -141,7 +183,6 @@ static bool build_compile_request(const KernelDefinition & definition, static std::shared_ptr load_kernel_executable(const KernelExecutablePrepareContext & context, const KernelDefinition & definition, const Dispatch & dispatch, - const std::vector & constants, const std::string & key, ggml_hrx_loom_jit_compile_result & compiled, std::string & error_message) { @@ -156,8 +197,8 @@ static std::shared_ptr load_kernel_executable(const KernelExec return nullptr; } - auto executable = std::make_shared(); - executable->launch = compiled.launch_config; + auto executable = std::make_shared(); + executable->launch_program = std::exchange(compiled.launch_program, nullptr); if (ErrorResult error = take_status(hrx_executable_load_data(context.device, compiled.hsaco_data, compiled.hsaco_size, "amdgpu", context.target, &executable->executable))) { @@ -177,18 +218,17 @@ static std::shared_ptr load_kernel_executable(const KernelExec GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); return nullptr; } - if (executable->export_info.binding_count != dispatch.bindings.size() || - executable->export_info.constant_byte_length != constants.size() || + const size_t launch_parameter_count = definition.launch_parameters.size(); + const bool constant_abi_valid = launch_parameter_count == 0 ? executable->export_info.constant_byte_length == 0 : + executable->export_info.constant_byte_length == launch_parameter_count * sizeof(uint32_t) || + executable->export_info.constant_byte_length == launch_parameter_count * sizeof(uint64_t); + if (executable->launch_program == nullptr || executable->export_info.binding_count != dispatch.bindings.size() || + !constant_abi_valid || executable->export_info.parameter_count != dispatch.bindings.size() + definition.launch_parameters.size()) { error_message = "compiled ABI does not match manifest for " + key; GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); return nullptr; } - if (executable->launch.workgroup_count[0] == 0 || executable->launch.workgroup_size[0] == 0) { - error_message = "compiled launch geometry is empty for " + key; - GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); - return nullptr; - } return executable; } @@ -224,7 +264,26 @@ class KernelExecutableCacheEntry { std::atomic load_state = LoadState::Unloaded; }; +static bool prepare_kernel_invocation(const KernelExecutableCacheEntry & entry, + const KernelExecutable & executable, + const KernelSpecialization & specialization, + std::vector & constants, + ggml_hrx_loom_jit_launch_config & launch) { + std::vector workload; + if (!build_workload_arguments(*entry.definition, specialization, workload)) { + return false; + } + if (ErrorResult error = take_status(ggml_hrx_loom_jit_launch_program_evaluate( + executable.launch_program, workload.empty() ? nullptr : workload.data(), workload.size(), &launch))) { + GGML_LOG_ERROR("%s: evaluate launch %s: %s\n", __func__, entry.key.c_str(), error->c_str()); + return false; + } + return pack_kernel_constants(*entry.definition, specialization, executable.export_info.constant_byte_length, + constants); +} + KernelExecutable::~KernelExecutable() { + ggml_hrx_loom_jit_launch_program_release(launch_program); if (executable != nullptr) { hrx_executable_release(executable); } @@ -265,12 +324,8 @@ bool KernelExecutableCache::ensure_jit_locked(const char * target, std::string & KernelExecutableRef KernelExecutableCache::get_or_compile(const KernelExecutablePrepareContext & context, const KernelDefinition & definition, - const Dispatch & dispatch, - std::vector & constants) { + const Dispatch & dispatch) { KernelExecutableRef ref; - if (!pack_kernel_constants(definition, dispatch, constants)) { - return ref; - } const std::string key = kernel_executable_key(definition, dispatch, context.target); LoomKernelCompileRequest request; @@ -300,7 +355,9 @@ KernelExecutableRef KernelExecutableCache::get_or_compile(const KernelExecutable std::shared_ptr KernelExecutableCache::materialize(const KernelExecutablePrepareContext & context, const KernelExecutableRef & ref, - const std::vector & constants) { + const KernelSpecialization & specialization, + std::vector & constants, + ggml_hrx_loom_jit_launch_config & launch) { if (!ref.valid()) { return nullptr; } @@ -308,7 +365,11 @@ std::shared_ptr KernelExecutableCache::materialize(const Kerne KernelExecutableCacheEntry & entry = *ref.entry; KernelExecutableCacheEntry::LoadState load_state = entry.load_state.load(std::memory_order_acquire); if (load_state == KernelExecutableCacheEntry::LoadState::Loaded) { - return std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + std::shared_ptr executable = + std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + return executable != nullptr && prepare_kernel_invocation(entry, *executable, specialization, constants, launch) + ? executable + : nullptr; } if (load_state == KernelExecutableCacheEntry::LoadState::Failed) { std::lock_guard entry_lock(entry.mutex); @@ -329,7 +390,11 @@ std::shared_ptr KernelExecutableCache::materialize(const Kerne GGML_LOG_ERROR("%s: %s\n", __func__, entry.error.c_str()); return nullptr; } - return std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + std::shared_ptr executable = + std::atomic_load_explicit(&entry.executable, std::memory_order_acquire); + return executable != nullptr && prepare_kernel_invocation(entry, *executable, specialization, constants, launch) + ? executable + : nullptr; } LoomCompiledKernelRef compiled_ref; @@ -360,7 +425,7 @@ std::shared_ptr KernelExecutableCache::materialize(const Kerne ggml_hrx_loom_jit_compile_result compiled = compiled_ref->take_result(); std::shared_ptr executable = load_kernel_executable( - context, *entry.definition, entry.dispatch, constants, entry.key, compiled, error_message); + context, *entry.definition, entry.dispatch, entry.key, compiled, error_message); if (executable != nullptr) { compiled.reset(); @@ -386,15 +451,19 @@ std::shared_ptr KernelExecutableCache::materialize(const Kerne cache_.erase(found); } } + if (executable != nullptr && !prepare_kernel_invocation(entry, *executable, specialization, constants, launch)) { + return nullptr; + } return executable; } std::shared_ptr KernelExecutableCache::prepare(const KernelExecutablePrepareContext & context, const KernelDefinition & definition, const Dispatch & dispatch, - std::vector & constants) { - const KernelExecutableRef ref = get_or_compile(context, definition, dispatch, constants); - return materialize(context, ref, constants); + std::vector & constants, + ggml_hrx_loom_jit_launch_config & launch) { + const KernelExecutableRef ref = get_or_compile(context, definition, dispatch); + return materialize(context, ref, dispatch.kernel, constants, launch); } void KernelExecutableCache::clear() { diff --git a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h index 2ae130d9918e..42d9c0674649 100644 --- a/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h +++ b/ggml/src/ggml-hrx/runtime/kernel-executable-cache.h @@ -22,7 +22,7 @@ struct KernelExecutable { hrx_executable_t executable = nullptr; uint32_t export_ordinal = 0; hrx_executable_export_info_t export_info = {}; - ggml_hrx_loom_jit_launch_config launch; + ggml_hrx_loom_jit_launch_program * launch_program = nullptr; }; struct KernelExecutablePrepareContext { @@ -44,17 +44,19 @@ class KernelExecutableCache { KernelExecutableRef get_or_compile(const KernelExecutablePrepareContext & context, const KernelDefinition & definition, - const Dispatch & dispatch, - std::vector & constants); + const Dispatch & dispatch); std::shared_ptr materialize(const KernelExecutablePrepareContext & context, const KernelExecutableRef & ref, - const std::vector & constants); + const KernelSpecialization & specialization, + std::vector & constants, + ggml_hrx_loom_jit_launch_config & launch); std::shared_ptr prepare(const KernelExecutablePrepareContext & context, const KernelDefinition & definition, const Dispatch & dispatch, - std::vector & constants); + std::vector & constants, + ggml_hrx_loom_jit_launch_config & launch); void clear(); diff --git a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp index 693497e985d5..d79a1fbf5cfe 100644 --- a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp +++ b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp @@ -90,9 +90,11 @@ static bool compile_kernel(ggml_hrx_loom_jit_amdgpu * jit, compile_options.dependency_count = request.dependencies.size(); compile_options.config_bindings = configs.data(); compile_options.config_binding_count = configs.size(); - compile_options.workload_arguments = request.workload.data(); - compile_options.workload_argument_count = request.workload.size(); - compile_options.evaluate_launch_config = true; + compile_options.specialize_workload = request.specialize_workload; + compile_options.workload_arguments = request.specialize_workload ? request.workload.data() : nullptr; + compile_options.workload_argument_count = request.specialize_workload ? request.workload.size() : 0; + compile_options.load_launch_config = true; + compile_options.evaluate_launch_config = request.specialize_workload; if (ErrorResult error = take_status(ggml_hrx_loom_jit_amdgpu_compile(jit, &compile_options, &compiled))) { error_message = "compile " + key + ": " + *error; diff --git a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h index 627c20d3f23a..3353297290fc 100644 --- a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h +++ b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.h @@ -24,6 +24,7 @@ struct LoomKernelCompileRequest { std::vector dependencies; std::vector> config_storage; std::vector workload; + bool specialize_workload = true; }; class LoomCompiledKernel { diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl.work.loom b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl.work.loom index b5fec8f8cd68..9f14062192b4 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl.work.loom +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl.work.loom @@ -42,7 +42,7 @@ kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(% kernel.decl @llm_gated_delta_net_f32_wmma_head128() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_inplace() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_inout: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_snapshot() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %rms_scales: buffer) -kernel.decl @llm_gated_delta_net_projection_epilogue_f32() launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +kernel.decl @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) kernel.decl @llm_ssm_conv_dconv4_silu_rollback_f32() launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) // Scenario: tg_c5 @@ -177,13 +177,14 @@ check.case public @qwen38_27b_udq4kxl_tg_c5_009_ggml_get_rows_f32_case { check.benchmark<@qwen38_27b_udq4kxl_tg_c5_009_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_009_ggml_get_rows_f32 check.case public @qwen38_27b_udq4kxl_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(96) : index %alpha_raw = check.generate.fill value(0) : tensor<384xi8> %beta_raw = check.generate.fill value(0) : tensor<384xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<384xi8> %beta_dst = check.generate.fill value(0) : tensor<384xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) check.return } @@ -518,13 +519,14 @@ check.case public @qwen38_27b_udq4kxl_tg_c5_037_llm_ssm_conv_dconv4_silu_rollbac check.benchmark<@qwen38_27b_udq4kxl_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(240) : index %alpha_raw = check.generate.fill value(0) : tensor<960xi8> %beta_raw = check.generate.fill value(0) : tensor<960xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<960xi8> %beta_dst = check.generate.fill value(0) : tensor<960xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) check.return } @@ -883,13 +885,14 @@ check.case public @qwen38_27b_udq4kxl_tg_c5_067_llm_ssm_conv_dconv4_silu_rollbac check.benchmark<@qwen38_27b_udq4kxl_tg_c5_067_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_tg_c5_067_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_tg_c5_068_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(768) : index %alpha_raw = check.generate.fill value(0) : tensor<3072xi8> %beta_raw = check.generate.fill value(0) : tensor<3072xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<3072xi8> %beta_dst = check.generate.fill value(0) : tensor<3072xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) check.return } @@ -1246,13 +1249,14 @@ check.case public @qwen38_27b_udq4kxl_tg_c5_095_llm_ssm_conv_dconv4_silu_rollbac check.benchmark<@qwen38_27b_udq4kxl_tg_c5_095_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_tg_c5_095_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_tg_c5_096_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(192) : index %alpha_raw = check.generate.fill value(0) : tensor<768xi8> %beta_raw = check.generate.fill value(0) : tensor<768xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<768xi8> %beta_dst = check.generate.fill value(0) : tensor<768xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) check.return } @@ -1894,13 +1898,14 @@ check.case public @qwen38_27b_udq4kxl_tg_c5_146_llm_ssm_conv_dconv4_silu_rollbac check.benchmark<@qwen38_27b_udq4kxl_tg_c5_146_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_tg_c5_146_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_tg_c5_147_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(48) : index %alpha_raw = check.generate.fill value(0) : tensor<192xi8> %beta_raw = check.generate.fill value(0) : tensor<192xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<192xi8> %beta_dst = check.generate.fill value(0) : tensor<192xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) check.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl_recurrent.work.loom b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl_recurrent.work.loom index e23036e44ac0..df898f2bde8d 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl_recurrent.work.loom +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl_recurrent.work.loom @@ -28,7 +28,7 @@ kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mu kernel.decl target(@ggml_rmsnorm_gfx11_wave32) @ggml_rmsnorm_mul_rope_f32() launch(%input: buffer, %weight: buffer, %positions: buffer, %output: buffer) kernel.decl target(@ggml_set_rows_gfx11_wave64) @ggml_set_rows(%token_count: index, %cache_row_count: index, %hidden_size: index) launch(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_snapshot() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %rms_scales: buffer) -kernel.decl @llm_gated_delta_net_projection_epilogue_f32() launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +kernel.decl @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) kernel.decl @llm_ssm_conv_dconv4_silu_rollback_f32() launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) // Scenario: tg_c5 @@ -143,13 +143,14 @@ check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_007_ggml_get_rows_f32_case check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_007_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_007_ggml_get_rows_f32 check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_008_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(96) : index %alpha_raw = check.generate.fill value(0) : tensor<384xi8> %beta_raw = check.generate.fill value(0) : tensor<384xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<384xi8> %beta_dst = check.generate.fill value(0) : tensor<384xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) check.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json index 7858d8974d68..d6f812883b17 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json @@ -363,13 +363,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_010_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "96", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 96, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 96 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -379,7 +380,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_011_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -1272,13 +1278,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_038_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "192", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 96, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 192 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -1288,7 +1295,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_039_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -2081,13 +2093,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_064_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "240", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 240 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -2097,7 +2110,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_065_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -2942,13 +2960,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_092_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "144", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 144 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -2958,7 +2977,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_093_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -3803,13 +3827,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_120_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "768", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 768 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -3819,7 +3844,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_121_llm_gated_delta_net_f32_wmma_head128", @@ -5541,13 +5571,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_175_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "48", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 48 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -5557,7 +5588,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_176_llm_gated_delta_net_f32_wmma_head128_inplace", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom index 6c3225e9b420..3703a6f02ae5 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom @@ -43,7 +43,7 @@ kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(% kernel.decl @llm_gated_delta_net_f32_wmma_head128() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_inplace() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_inout: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_snapshot() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %rms_scales: buffer) -kernel.decl @llm_gated_delta_net_projection_epilogue_f32() launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +kernel.decl @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) kernel.decl @llm_ssm_conv_dconv4_silu_rollback_f32() launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) @@ -181,13 +181,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_009_ggml_get_rows_f32_case { check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_009_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_009_ggml_get_rows_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_010_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(96) : index %alpha_raw = check.generate.fill value(0) : tensor<384xi8> %beta_raw = check.generate.fill value(0) : tensor<384xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<384xi8> %beta_dst = check.generate.fill value(0) : tensor<384xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) check.return } @@ -545,13 +546,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_037_llm_ssm_conv_dconv4_silu_r check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_037_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_037_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_038_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(192) : index %alpha_raw = check.generate.fill value(0) : tensor<768xi8> %beta_raw = check.generate.fill value(0) : tensor<768xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<768xi8> %beta_dst = check.generate.fill value(0) : tensor<768xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) check.return } @@ -888,13 +890,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_063_llm_ssm_conv_dconv4_silu_r check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_063_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_063_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_064_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(240) : index %alpha_raw = check.generate.fill value(0) : tensor<960xi8> %beta_raw = check.generate.fill value(0) : tensor<960xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<960xi8> %beta_dst = check.generate.fill value(0) : tensor<960xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) check.return } @@ -1252,13 +1255,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_091_llm_ssm_conv_dconv4_silu_r check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_091_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_091_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_092_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(144) : index %alpha_raw = check.generate.fill value(0) : tensor<576xi8> %beta_raw = check.generate.fill value(0) : tensor<576xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<576xi8> %beta_dst = check.generate.fill value(0) : tensor<576xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<576xi8>, tensor<576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<576xi8>, tensor<576xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<576xi8>, tensor<576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<576xi8>, tensor<576xi8>) check.return } @@ -1616,13 +1620,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_119_llm_ssm_conv_dconv4_silu_r check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_119_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_119_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_120_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(768) : index %alpha_raw = check.generate.fill value(0) : tensor<3072xi8> %beta_raw = check.generate.fill value(0) : tensor<3072xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<3072xi8> %beta_dst = check.generate.fill value(0) : tensor<3072xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) check.return } @@ -2310,13 +2315,14 @@ check.case public @qwen38_27b_udq4kxl_draft_tg_c1_174_llm_ssm_conv_dconv4_silu_r check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_174_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_174_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_draft_tg_c1_175_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(48) : index %alpha_raw = check.generate.fill value(0) : tensor<192xi8> %beta_raw = check.generate.fill value(0) : tensor<192xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<192xi8> %beta_dst = check.generate.fill value(0) : tensor<192xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) check.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json index 2e72c85c3677..c25c23af3c4d 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json @@ -363,13 +363,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "96", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 96, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 96 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -379,7 +380,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -1272,13 +1278,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "192", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 96, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 192 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -1288,7 +1295,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -2081,13 +2093,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_064_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "240", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 240 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -2097,7 +2110,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_065_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -2942,13 +2960,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_092_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "144", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 144 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -2958,7 +2977,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_093_llm_gated_delta_net_f32_wmma_head128_snapshot", @@ -3803,13 +3827,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_120_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "768", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 768 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -3819,7 +3844,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_121_llm_gated_delta_net_f32_wmma_head128", @@ -5561,13 +5591,14 @@ { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_176_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "48", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 48 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -5577,7 +5608,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_177_llm_gated_delta_net_f32_wmma_head128_inplace", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.work.loom b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.work.loom index 841905c9608e..e92702b54836 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.work.loom +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.work.loom @@ -43,7 +43,7 @@ kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(% kernel.decl @llm_gated_delta_net_f32_wmma_head128() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_inplace() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_inout: buffer, %dst: buffer, %rms_scales: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_snapshot() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %rms_scales: buffer) -kernel.decl @llm_gated_delta_net_projection_epilogue_f32() launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +kernel.decl @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) kernel.decl @llm_ssm_conv_dconv4_silu_rollback_f32() launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) @@ -181,13 +181,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_009_ggml_get_rows_f32_case check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_009_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_009_ggml_get_rows_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(96) : index %alpha_raw = check.generate.fill value(0) : tensor<384xi8> %beta_raw = check.generate.fill value(0) : tensor<384xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<384xi8> %beta_dst = check.generate.fill value(0) : tensor<384xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<384xi8>, tensor<384xi8>, tensor<192xi8>, tensor<192xi8>, tensor<384xi8>, tensor<384xi8>) check.return } @@ -545,13 +546,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_037_llm_ssm_conv_dconv4_si check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(192) : index %alpha_raw = check.generate.fill value(0) : tensor<768xi8> %beta_raw = check.generate.fill value(0) : tensor<768xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<768xi8> %beta_dst = check.generate.fill value(0) : tensor<768xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) check.return } @@ -888,13 +890,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_063_llm_ssm_conv_dconv4_si check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_063_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_063_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_064_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(240) : index %alpha_raw = check.generate.fill value(0) : tensor<960xi8> %beta_raw = check.generate.fill value(0) : tensor<960xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<960xi8> %beta_dst = check.generate.fill value(0) : tensor<960xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<960xi8>, tensor<960xi8>, tensor<192xi8>, tensor<192xi8>, tensor<960xi8>, tensor<960xi8>) check.return } @@ -1252,13 +1255,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_091_llm_ssm_conv_dconv4_si check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_091_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_091_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_092_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(144) : index %alpha_raw = check.generate.fill value(0) : tensor<576xi8> %beta_raw = check.generate.fill value(0) : tensor<576xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<576xi8> %beta_dst = check.generate.fill value(0) : tensor<576xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<576xi8>, tensor<576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<576xi8>, tensor<576xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<576xi8>, tensor<576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<576xi8>, tensor<576xi8>) check.return } @@ -1616,13 +1620,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_119_llm_ssm_conv_dconv4_si check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_119_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_119_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_120_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(768) : index %alpha_raw = check.generate.fill value(0) : tensor<3072xi8> %beta_raw = check.generate.fill value(0) : tensor<3072xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<3072xi8> %beta_dst = check.generate.fill value(0) : tensor<3072xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<3072xi8>, tensor<3072xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3072xi8>, tensor<3072xi8>) check.return } @@ -2319,13 +2324,14 @@ check.case public @qwen38_27b_udq4kxl_companion_tg_c5_175_llm_ssm_conv_dconv4_si check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_175_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_175_llm_ssm_conv_dconv4_silu_rollback_f32 check.case public @qwen38_27b_udq4kxl_companion_tg_c5_176_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(48) : index %alpha_raw = check.generate.fill value(0) : tensor<192xi8> %beta_raw = check.generate.fill value(0) : tensor<192xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<192xi8> %beta_dst = check.generate.fill value(0) : tensor<192xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>, tensor<192xi8>) check.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json index afeb1ccedd7e..5fb4222c0d2b 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json @@ -309,13 +309,14 @@ { "benchmark": "@v2_mtp4_program4_tg_c4_008_llm_gated_delta_net_projection_epilogue_f32", "compile_parameters": { - "llm.gated_delta_net.epilogue_element_count": "192", "llm.gated_delta_net.epilogue_head_count": "48", "llm.gated_delta_net.epilogue_workgroup_size": "256" }, "corpus_dir": "kernel-corpus/kernels/loom-libs", "count": 48, - "integer_parameters": {}, + "integer_parameters": { + "element_count": 192 + }, "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", "library_sources": [], "primary_sources": [ @@ -325,7 +326,12 @@ "ops/gated_delta_net_f32_wmma.loom" ], "symbol": "llm_gated_delta_net_projection_epilogue_f32", - "workload_parameters": [] + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] }, { "benchmark": "@v2_mtp4_program4_tg_c4_009_llm_gated_delta_net_f32_wmma_head128_snapshot", diff --git a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.work.loom b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.work.loom index 3b3e610abcbc..9f5922a7a8b3 100644 --- a/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.work.loom +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.work.loom @@ -26,7 +26,7 @@ kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mu kernel.decl target(@ggml_rmsnorm_gfx11_wave32) @ggml_rmsnorm_mul_rope_f32() launch(%input: buffer, %weight: buffer, %positions: buffer, %output: buffer) kernel.decl target(@ggml_set_rows_gfx11_wave64) @ggml_set_rows(%token_count: index, %cache_row_count: index, %hidden_size: index) launch(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) kernel.decl @llm_gated_delta_net_f32_wmma_head128_snapshot() launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %rms_scales: buffer) -kernel.decl @llm_gated_delta_net_projection_epilogue_f32() launch(%alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) +kernel.decl @llm_gated_delta_net_projection_epilogue_f32(%element_count: index) launch(%element_count: index, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %gate_dst: buffer, %beta_dst: buffer) kernel.decl @llm_ssm_conv_dconv4_silu_rollback_f32() launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) @@ -144,13 +144,14 @@ check.case public @v2_mtp4_program4_tg_c4_007_ggml_get_rows_f32_case { check.benchmark<@v2_mtp4_program4_tg_c4_007_ggml_get_rows_f32_case> @v2_mtp4_program4_tg_c4_007_ggml_get_rows_f32 check.case public @v2_mtp4_program4_tg_c4_008_llm_gated_delta_net_projection_epilogue_f32_case { + %element_count = check.literal value(192) : index %alpha_raw = check.generate.fill value(0) : tensor<768xi8> %beta_raw = check.generate.fill value(0) : tensor<768xi8> %bias = check.generate.fill value(0) : tensor<192xi8> %a_scale = check.generate.fill value(0) : tensor<192xi8> %gate_dst = check.generate.fill value(0) : tensor<768xi8> %beta_dst = check.generate.fill value(0) : tensor<768xi8> - kernel.launch @llm_gated_delta_net_projection_epilogue_f32[](%alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [](tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) + kernel.launch @llm_gated_delta_net_projection_epilogue_f32[%element_count](%element_count, %alpha_raw, %beta_raw, %bias, %a_scale, %gate_dst, %beta_dst) : [index](index, tensor<768xi8>, tensor<768xi8>, tensor<192xi8>, tensor<192xi8>, tensor<768xi8>, tensor<768xi8>) check.return } diff --git a/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py index e93268dc2683..551fa6fe8fb9 100755 --- a/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py +++ b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py @@ -54,6 +54,27 @@ def source_path(entry: dict[str, Any], source: str) -> Path: return HRX_DIR / corpus_dir / source +def loom_link_source_args(entry: dict[str, Any]) -> list[str]: + sources = entry.get("sources") + if sources is None: + sources = [*entry.get("primary_sources", []), *entry.get("library_sources", [])] + return [str(source_path(entry, source)) for source in sources] + + +def loom_link_provider_source_args(entry: dict[str, Any]) -> list[str]: + primary_sources = entry.get("primary_sources") + library_sources = entry.get("library_sources") + closure = [str(source_path(entry, source)) for source in library_sources or []] + providers = [f"--library={source_path(entry, source)}" for source in primary_sources or []] + return [*closure, *providers] + + +def missed_exact_source_selection(result: CommandResult) -> bool: + if result.state == "ok" or not result.stderr.is_file(): + return False + return "exact link selection missed reachable source symbol ref" in result.stderr.read_text(encoding="utf-8") + + def manifest_entries(manifest: dict[str, Any]) -> list[dict[str, Any]]: if "dispatches" in manifest: return manifest.get("dispatches", []) @@ -384,10 +405,7 @@ def main() -> None: stage_results: dict[str, Any] = {} config_args = [f"--config={name}={value}" for name, value in sorted(entry.get("compile_parameters", {}).items())] - sources = entry.get("sources") - if sources is None: - sources = [*entry.get("primary_sources", []), *entry.get("library_sources", [])] - link_inputs = [str(source_path(entry, source)) for source in sources] + link_inputs = loom_link_source_args(entry) state = "ok" error = None link_result = run_command( @@ -407,6 +425,25 @@ def main() -> None: "loom-link", ) stage_results["loom-link"] = command_result_json(link_result) + if missed_exact_source_selection(link_result) and entry.get("primary_sources"): + provider_link_inputs = loom_link_provider_source_args(entry) + link_result = run_command( + [ + str(loom_link), + str(loom_source), + *provider_link_inputs, + "--mode=link", + "--to=text", + f"--root={benchmark}", + f"--output={linked_source}", + *config_args, + ], + REPO_DIR, + args.benchmark_timeout_sec, + bench_dir, + "loom-link-provider-retry", + ) + stage_results["loom-link-provider-retry"] = command_result_json(link_result) state = link_result.state error = link_result.error if state == "ok": diff --git a/ggml/src/ggml-hrx/tools/benchmarks/test-run-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/test-run-model-benchmarks.py new file mode 100644 index 000000000000..97d33a75ddcb --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/test-run-model-benchmarks.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("run-model-benchmarks.py") +SPEC = importlib.util.spec_from_file_location("run_model_benchmarks", SCRIPT) +if SPEC is None or SPEC.loader is None: + raise RuntimeError(f"could not load {SCRIPT}") +MODULE = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = MODULE +SPEC.loader.exec_module(MODULE) + + +class LoomLinkSourceArgsTest(unittest.TestCase): + def test_preserves_provider_and_source_closure_roles_for_retry(self) -> None: + entry = { + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "sources": ["ops/example.loom", "motifs/helper.loom"], + "primary_sources": ["ops/example.loom"], + "library_sources": ["motifs/helper.loom"], + } + + args = MODULE.loom_link_provider_source_args(entry) + + corpus = MODULE.HRX_DIR / entry["corpus_dir"] + self.assertEqual( + args, + [ + str(corpus / "motifs/helper.loom"), + f"--library={corpus / 'ops/example.loom'}", + ], + ) + + def test_supports_legacy_flat_source_manifests(self) -> None: + entry = { + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "sources": ["ops/example.loom", "motifs/helper.loom"], + } + + args = MODULE.loom_link_source_args(entry) + + corpus = MODULE.HRX_DIR / entry["corpus_dir"] + self.assertEqual( + args, + [str(corpus / "ops/example.loom"), str(corpus / "motifs/helper.loom")], + ) + + def test_initial_link_retains_flat_source_order(self) -> None: + entry = { + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "sources": ["ops/example.loom", "motifs/helper.loom"], + "primary_sources": ["ops/example.loom"], + "library_sources": ["motifs/helper.loom"], + } + + args = MODULE.loom_link_source_args(entry) + + corpus = MODULE.HRX_DIR / entry["corpus_dir"] + self.assertEqual( + args, + [str(corpus / "ops/example.loom"), str(corpus / "motifs/helper.loom")], + ) + + def test_retries_only_missed_exact_source_selection(self) -> None: + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + stderr = root / "stderr.txt" + stderr.write_text( + "exact link selection missed reachable source symbol ref {module=0, symbol=2}\n", + encoding="utf-8", + ) + result = MODULE.CommandResult( + "failed", + "loom-link exited with status 1", + 1, + None, + root / "stdout.txt", + stderr, + root / "command.json", + ) + self.assertTrue(MODULE.missed_exact_source_selection(result)) + + stderr.write_text("unresolved exact declaration\n", encoding="utf-8") + self.assertFalse(MODULE.missed_exact_source_selection(result)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index db113e6e0970..2c2f930b7f10 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -6,6 +6,7 @@ #include "dispatch/command-program.h" #include "dispatch/dispatch-scheduler.h" #include "dispatch_registration/common/dispatch-gather-add.h" +#include "dispatch_registration/common/dispatch-symmetric-i4.h" #include "dispatch_registration/dispatch-registry.h" #include "ggml-alloc.h" #include "ggml-backend.h" @@ -281,6 +282,10 @@ static int64_t matmul_weight_format_config(ggml_type type) { return 50; case GGML_TYPE_Q5_1: return 51; + case GGML_TYPE_IQ1_S: + return 19; + case GGML_TYPE_IQ1_M: + return 29; case GGML_TYPE_IQ2_S: return 22; case GGML_TYPE_IQ3_S: @@ -1921,6 +1926,86 @@ static void run_binary_f32_broadcast_dispatch_checks() { } } +static void run_binary_q8_publication_checks() { + const struct { + int64_t hidden; + int64_t tokens; + bool broadcast; + bool multiply; + bool side_use; + int projections; + const char * expected_kernel; + } cases[] = { + { 2048, 2, false, true, false, 1, "loom_libs:ggml_binary_f32_publish_q8_1_x4" }, + { 2048, 4, false, false, true, 2, "loom_libs:ggml_binary_f32_publish_q8_1_x4" }, + { 256, 2, true, true, false, 1, "loom_libs:ggml_binary_bc_f32_publish_q8_1_x4" }, + { 2048, 6, false, true, false, 1, "loom_libs:ggml_binary_f32" }, + { 2048, 2, false, true, false, 0, "loom_libs:ggml_binary_f32" }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); + ggml_tensor * rhs = test.broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden) : + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); + ggml_tensor * binary = test.multiply ? ggml_mul(ctx, lhs, rhs) : ggml_add(ctx, lhs, rhs); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + REQUIRE(binary != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + for (int i = 0; i < test.projections; ++i) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.hidden, 512); + ggml_build_forward_expand(graph, ggml_mul_mat(ctx, weight, binary)); + } + if (test.projections == 0 || test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, binary, 0.5f)); + } + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const ggml::hrx::Value * binary_value = imported.graph.values().find_tensor(binary); + REQUIRE(binary_value != nullptr); + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return dispatch.bindings.size() >= 3 && dispatch.bindings[2].value == binary_value->id; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(kernel_name_for_id(producer->kernel.kernel_id) == test.expected_kernel); + + const bool publishes = std::string_view(test.expected_kernel).find("publish_q8_1_x4") != std::string_view::npos; + const size_t q8_bytes = qwen_q8_1_x4_size(test.tokens, test.hidden); + const auto * q8 = plan.metadata.find_alternate_value(binary_value->id, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE((q8 != nullptr) == publishes); + if (publishes) { + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings.back().value == q8->alternate_value); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; + }) == 0); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == q8->alternate_value; + }) == test.projections); + } + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + static void run_graph_snapshot_diagnostics_checks() { ggml_init_params params = {}; params.mem_size = 256 * 1024; @@ -3082,6 +3167,8 @@ static void run_graph_traversal_checks() { } } +static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output); + static void schedule_single_matmul_command(ggml_context * ctx, ggml_tensor * output, const char * expected_kernel_name, @@ -3106,6 +3193,10 @@ 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 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_") || + string_contains(kernel_name, "_iq3_xxs_"); const std::string expected_kernel(expected_kernel_name); const bool q8_capable_kernel = expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || expected_kernel == "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; @@ -3130,7 +3221,7 @@ 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() == (q6_vector_kernel ? 3 : 4)); + REQUIRE(dispatch.bindings.size() == (codebook_kernel ? 5 : q6_vector_kernel ? 3 : 4)); if (!q6_vector_kernel) { require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); } @@ -3159,8 +3250,13 @@ static void schedule_single_matmul_command(ggml_context * ctx, require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.weight_format", std::to_string(matmul_weight_format_config(weight->type))); } else if (string_contains(kernel_name, "ggml_mul_mat")) { - REQUIRE(dispatch.bindings.size() == 3); - require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + REQUIRE(dispatch.bindings.size() == (codebook_kernel ? 4 : 3)); + if (aligned_dynamic_kernel) { + REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); + } else { + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + } require_compile_parameter(dispatch, "ggml.mul_mat.input_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "ggml.mul_mat.output_size", std::to_string(expected_output_size)); require_compile_parameter(dispatch, "ggml.mul_mat.output_accumulation", "0"); @@ -3191,13 +3287,108 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == command_count); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.back().bindings.size() == (vector_kernel && !q6_vector_kernel ? 4 : 3)); + REQUIRE(commands.commands.back().bindings.size() == + (codebook_kernel ? (vector_kernel ? 5 : 4) : vector_kernel && !q6_vector_kernel ? 4 : 3)); REQUIRE(commands.commands.back().bindings[0].name == "input"); REQUIRE(commands.commands.back().bindings[1].name == "weight"); - REQUIRE(commands.commands.back().bindings[2].name == "output"); + const size_t output_binding = codebook_kernel ? 3 : 2; + if (codebook_kernel) { + REQUIRE(string_contains(commands.commands.back().bindings[2].name, "grid")); + } + REQUIRE(commands.commands.back().bindings[output_binding].name == "output"); if (vector_kernel && !q6_vector_kernel) { - REQUIRE(commands.commands.back().bindings[3].name == "next_output"); + REQUIRE(commands.commands.back().bindings[output_binding + 1].name == "next_output"); + } +} + +static void run_iq1_codebook_matmul_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + auto check_single = [&](ggml_type type, int64_t token_count, const char * expected_kernel) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches[0].kernel.kernel_id) == expected_kernel); + REQUIRE(scheduler.plan().transients.size() == 1); + REQUIRE(scheduler.plan().transients[0].name == "common.iq1.grid.v1"); + REQUIRE(scheduler.plan().transients[0].size == 16384); + REQUIRE(scheduler.plan().constant_initializations.size() == 1); + REQUIRE(scheduler.plan().constant_initializations[0].name == "common.iq1.grid.v1"); + REQUIRE(scheduler.plan().constant_initializations[0].data.size() == 16384); + const std::array expected_grid_1542 = { 0, 1, 0, 0xff, 1, 0, 0xff, 1 }; + REQUIRE(std::equal(expected_grid_1542.begin(), expected_grid_1542.end(), + scheduler.plan().constant_initializations[0].data.begin() + 1542 * 8)); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings.size() == (token_count == 1 ? 5 : 4)); + REQUIRE(commands.commands[0].bindings[0].name == "input"); + REQUIRE(commands.commands[0].bindings[1].name == "weight"); + REQUIRE(commands.commands[0].bindings[2].name == "iq1_grid"); + REQUIRE(commands.commands[0].bindings[3].name == "output"); + }; + + check_single(GGML_TYPE_IQ1_S, 1, "loom_libs:ggml_mul_mat_vector_iq1_s_f32_f32"); + check_single(GGML_TYPE_IQ1_S, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32"); + check_single(GGML_TYPE_IQ1_M, 1, "loom_libs:ggml_mul_mat_vector_iq1_m_f32_f32"); + check_single(GGML_TYPE_IQ1_M, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32"); + + for (const auto [type, disable_env] : + { std::pair{ GGML_TYPE_IQ1_S, "GGML_HRX_DISABLE_IQ1_S_CODEBOOK_MATMUL" }, + std::pair{ GGML_TYPE_IQ1_M, "GGML_HRX_DISABLE_IQ1_M_CODEBOOK_MATMUL" } }) { + EnvironmentVariableGuard guard(disable_env); + REQUIRE(setenv(disable_env, "1", 1) == 0); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, output)); } + + ggml_tensor * iq1_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ1_S, 2048, 128); + ggml_tensor * iq1_m_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ1_M, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * iq1_s_output = ggml_mul_mat(ctx, iq1_s_weight, input); + ggml_tensor * iq1_m_output = ggml_mul_mat(ctx, iq1_m_weight, input); + REQUIRE(iq1_s_output != nullptr); + REQUIRE(iq1_m_output != nullptr); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, iq1_s_output); + ggml_build_forward_expand(graph, iq1_m_output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + REQUIRE(scheduler.plan().transients.size() == 1); + REQUIRE(scheduler.plan().transients[0].name == "common.iq1.grid.v1"); + REQUIRE(scheduler.plan().constant_initializations.size() == 1); + REQUIRE(scheduler.plan().dispatches[0].bindings[2].value == scheduler.plan().transients[0].value); + REQUIRE(scheduler.plan().dispatches[1].bindings[2].value == scheduler.plan().transients[0].value); + + ggml_free(ctx); } static void schedule_fused_matmul_unary_command(ggml_context * ctx, @@ -3749,6 +3940,181 @@ static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { ggml_free(ctx); } +static void run_qwen_decode_rmsnorm_publication_checks() { + const struct { + bool projection; + bool side_use; + bool q8; + } cases[] = { + { true, false, true }, + { true, true, true }, + { false, false, false }, + }; + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + ggml_tensor * normalized = ggml_mul(ctx, ggml_rms_norm(ctx, input, 0.000001f), norm_weight); + ggml_tensor * projection = nullptr; + if (test.projection) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 512); + projection = ggml_mul_mat(ctx, weight, normalized); + } + ggml_tensor * side = test.side_use ? ggml_scale(ctx, normalized, 0.5f) : nullptr; + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, projection != nullptr ? projection : normalized); + if (side != nullptr) { + ggml_build_forward_expand(graph, side); + } + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"; + }); + REQUIRE((producer != plan.dispatches.end()) == test.q8); + const ggml::hrx::Value * output = imported.graph.values().find_tensor(normalized); + REQUIRE(output != nullptr); + const auto * q8 = plan.metadata.find_alternate_value( + output->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(1, 2048)); + REQUIRE((q8 != nullptr) == test.q8); + if (test.q8) { + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings.back().value == q8->alternate_value); + REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { + return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; + })); + } + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + +static void run_qwen_decode_publication_consumer_qualification_checks() { + enum class ConsumerCase { + Dense, + UnsupportedWeight, + WrongOperand, + IncompatibleGeometry, + Routed, + RoutedQ6, + RoutedShape, + RoutedRoute, + AliasedDense, + }; + const struct { + ConsumerCase consumer; + bool q8; + } cases[] = { + { ConsumerCase::Dense, true }, + { ConsumerCase::UnsupportedWeight, false }, + { ConsumerCase::WrongOperand, false }, + { ConsumerCase::IncompatibleGeometry, false }, + { ConsumerCase::Routed, true }, + { ConsumerCase::RoutedQ6, true }, + { ConsumerCase::RoutedShape, false }, + { ConsumerCase::RoutedRoute, false }, + { ConsumerCase::AliasedDense, true }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); + REQUIRE(rms != nullptr); + REQUIRE(normalized != nullptr); + + ggml_tensor * consumer_input = normalized; + if (test.consumer == ConsumerCase::AliasedDense) { + consumer_input = ggml_reshape_2d(ctx, normalized, 2048, 1); + REQUIRE(consumer_input != nullptr); + } + + ggml_tensor * terminal = nullptr; + if (test.consumer == ConsumerCase::Dense || test.consumer == ConsumerCase::AliasedDense) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 512); + terminal = ggml_mul_mat(ctx, weight, consumer_input); + } else if (test.consumer == ConsumerCase::UnsupportedWeight) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 512); + terminal = ggml_mul_mat(ctx, weight, normalized); + } else if (test.consumer == ConsumerCase::WrongOperand) { + ggml_tensor * activation = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + terminal = ggml_mul_mat(ctx, normalized, activation); + } else if (test.consumer == ConsumerCase::IncompatibleGeometry) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 63); + terminal = ggml_mul_mat(ctx, weight, normalized); + } else { + const int64_t route_count = test.consumer == ConsumerCase::RoutedShape ? 4 : 8; + ggml_tensor * route_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, 1); + ggml_tensor * gate_weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, 128); + ggml_tensor * up_weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, 128); + const ggml_type down_weight_type = test.consumer == ConsumerCase::RoutedQ6 ? GGML_TYPE_Q6_K : GGML_TYPE_Q4_K; + ggml_tensor * down_weight = ggml_new_tensor_3d(ctx, down_weight_type, 768, 2048, 128); + ggml_tensor * gate = ggml_mul_mat_id(ctx, gate_weight, normalized, route_ids); + ggml_tensor * up = ggml_mul_mat_id(ctx, up_weight, normalized, route_ids); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + ggml_tensor * down_route_ids = test.consumer == ConsumerCase::RoutedRoute ? + ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, 1) : + route_ids; + ggml_tensor * down = ggml_mul_mat_id(ctx, down_weight, glu, down_route_ids); + ggml_tensor * route_weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, 1); + terminal = ggml_mul(ctx, down, route_weights); + } + REQUIRE(terminal != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, terminal); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + const ggml::hrx::Value * rms_value = imported.graph.values().find_tensor(rms); + const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); + REQUIRE(rms_value != nullptr); + REQUIRE(normalized_value != nullptr); + const ggml::hrx::GraphNode * rms_node = imported.graph.index().producer(rms_value->id); + size_t rms_index = 0; + REQUIRE(rms_node != nullptr); + REQUIRE(imported.graph.index().node_index(rms_node, rms_index)); + + const ggml::hrx::CommandPlan empty_plan; + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, rms_index, match)); + const auto * q8 = match.metadata.find_alternate_value( + normalized_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(1, 2048)); + REQUIRE((q8 != nullptr) == test.q8); + if (test.q8) { + REQUIRE(kernel_name_for_id(match.dispatches.back().kernel.kernel_id) == + "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); + } + ggml_free(ctx); + } +} + static void schedule_qwen_terminal_q6k_vector_command() { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -4175,10 +4541,12 @@ static void run_get_rows_rmsnorm_binary_dispatch_checks() { const size_t f16_bytes = output_value->byte_count / 2; const auto * q8 = plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_bytes); const auto * f16 = plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_F16, f16_bytes); - REQUIRE(q8 != nullptr); - REQUIRE(f16 != nullptr); - REQUIRE(dispatch.bindings[5].value == q8->alternate_value); - REQUIRE(dispatch.bindings[6].value == f16->alternate_value); + REQUIRE(q8 == nullptr); + REQUIRE(f16 == nullptr); + REQUIRE(dispatch.bindings[5].value == plan.transients[0].value); + REQUIRE(dispatch.bindings[6].value == plan.transients[1].value); + REQUIRE(plan.transients[0].size == q8_bytes); + REQUIRE(plan.transients[1].size == f16_bytes); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -4776,6 +5144,9 @@ static void schedule_flash_attention_command(ggml_context * ctx, REQUIRE(scheduler.plan().transients.size() == 1); const auto & copy = scheduler.plan().dispatches.front(); REQUIRE(kernel_name_for_id(copy.kernel.kernel_id) == "loom_libs:ggml_copy_transpose_f16"); + REQUIRE(copy.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(copy.kernel.integer_parameters.at("row_count") == key_value_token_count); + REQUIRE(copy.kernel.integer_parameters.at("column_count") == key_value_head_count * value_head_size); REQUIRE(copy.bindings.size() == 2); REQUIRE(copy.bindings[1].value == dispatch.bindings[2].value); REQUIRE(copy.bindings[0].value != copy.bindings[1].value); @@ -4820,16 +5191,15 @@ static void schedule_common_flash_attention_command(ggml_context * ctx, query_head_count, key_value_head_count, qk_head_size, value_head_size, scale); } -static void schedule_common_flash_attention_decode_split_command(ggml_context * ctx, - ggml_tensor * output, - int64_t query_token_count, - int64_t key_value_token_count, - int64_t query_head_count, - int64_t key_value_head_count, - int64_t qk_head_size = kQwenFlashHeadSize, - int64_t value_head_size = kQwenFlashHeadSize, - float scale = 1.0f / std::sqrt(128.0f)) { - const int64_t key_value_capacity = ((key_value_token_count + 63) / 64) * 64; +static void schedule_common_flash_attention_decode_fallback_command(ggml_context * ctx, + ggml_tensor * output, + int64_t query_token_count, + int64_t key_value_token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t qk_head_size = kQwenFlashHeadSize, + int64_t value_head_size = kQwenFlashHeadSize, + float scale = 1.0f / std::sqrt(128.0f)) { ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); @@ -4843,21 +5213,22 @@ static void schedule_common_flash_attention_decode_split_command(ggml_context * REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == static_cast(query_token_count)); + REQUIRE(scheduler.plan().transients.size() == 4); + REQUIRE(scheduler.plan().completion_counter_requests.size() == 1); + REQUIRE(scheduler.plan().completion_counter_requests.front().count == key_value_head_count); + REQUIRE(scheduler.plan().metadata.alternate_values().empty()); - for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); - REQUIRE(dispatch.bindings.size() == 10); - require_compile_parameter(dispatch, "ggml.flash_attention.query_head_count", std::to_string(query_head_count)); - require_compile_parameter(dispatch, "ggml.flash_attention.key_value_head_count", - std::to_string(key_value_head_count)); - require_compile_parameter(dispatch, "ggml.flash_attention.qk_head_size", std::to_string(qk_head_size)); - require_compile_parameter(dispatch, "ggml.flash_attention.value_head_size", std::to_string(value_head_size)); - require_compile_parameter(dispatch, "ggml.flash_attention.attention_scale", expected_config_value(scale)); - require_compile_parameter(dispatch, "ggml.flash_attention.decode.key_value_token_capacity", - std::to_string(key_value_capacity)); - } + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); + REQUIRE(dispatch.bindings.size() == 10); + require_compile_parameter(dispatch, "ggml.flash_attention.query_head_count", std::to_string(query_head_count)); + require_compile_parameter(dispatch, "ggml.flash_attention.key_value_head_count", + std::to_string(key_value_head_count)); + require_compile_parameter(dispatch, "ggml.flash_attention.qk_head_size", std::to_string(qk_head_size)); + require_compile_parameter(dispatch, "ggml.flash_attention.value_head_size", std::to_string(value_head_size)); + require_compile_parameter(dispatch, "ggml.flash_attention.attention_scale", expected_config_value(scale)); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -4869,14 +5240,105 @@ static void schedule_common_flash_attention_decode_split_command(ggml_context * REQUIRE(commands.commands.front().bindings[1].name == "key"); REQUIRE(commands.commands.front().bindings[2].name == "value"); REQUIRE(commands.commands.front().bindings[3].name == "mask"); - REQUIRE(commands.commands.front().bindings[4].name == "partial_max"); - REQUIRE(commands.commands.front().bindings[5].name == "partial_sum"); - REQUIRE(commands.commands.front().bindings[6].name == "partial_output"); - REQUIRE(commands.commands.front().bindings[7].name == "completion_counter"); REQUIRE(commands.commands.front().bindings[8].name == "output"); REQUIRE(commands.commands.front().bindings[9].name == "next_q8_output"); } +static void schedule_common_flash_attention_decode_q8_command(ggml_context * ctx, ggml_type weight_type) { + constexpr int64_t query_token_count = 1; + constexpr int64_t key_value_token_count = 512; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t hidden_size = query_head_count * kQwenFlashHeadSize; + constexpr int64_t output_size = 256; + const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, hidden_size); + + ggml_tensor * output = build_qwen_flash_attention_graph( + ctx, query_token_count, key_value_token_count, query_head_count, key_value_head_count); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden_size, query_token_count); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, hidden_size, output_size); + ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); + REQUIRE(reshaped != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(projected != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, projected); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 2); + REQUIRE(plan.transients.size() == 4); + REQUIRE(plan.completion_counter_requests.size() == 1); + + const ggml::hrx::Value * native_value = imported.graph.values().find_tensor(output); + const ggml::hrx::Value * reshape_value = imported.graph.values().find_tensor(reshaped); + REQUIRE(native_value != nullptr); + REQUIRE(reshape_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * q8 = + ggml::hrx::find_alternate_value(plan, reshape_value->id, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(q8 != nullptr); + REQUIRE(ggml::hrx::find_alternate_value(plan, native_value->id, GGML_TYPE_Q8_1, q8_bytes) == nullptr); + + const ggml::hrx::Dispatch & producer = plan.dispatches.front(); + const ggml::hrx::Dispatch & consumer = plan.dispatches.back(); + REQUIRE(kernel_name_for_id(producer.kernel.kernel_id) == + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + REQUIRE(producer.bindings.size() == 10); + REQUIRE(producer.bindings[9].value == q8->alternate_value); + REQUIRE(consumer.bindings[0].value == q8->alternate_value); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 2); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 10); + REQUIRE(commands.commands.front().bindings[8].name == "output"); + REQUIRE(commands.commands.front().bindings[9].name == "next_q8_output"); +} + +static void schedule_common_flash_attention_decode_unqualified_command(ggml_context * ctx) { + constexpr int64_t hidden_size = 4 * kQwenFlashHeadSize; + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 512, 4, 2); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden_size, 1); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 256); + ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); + REQUIRE(reshaped != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(projected != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, projected); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.metadata.alternate_values().empty()); + REQUIRE(plan.dispatches.size() == 2); + REQUIRE(plan.transients.size() == 4); + REQUIRE(plan.completion_counter_requests.size() == 1); + REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + REQUIRE(plan.dispatches.front().bindings.size() == 10); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.size() == 2); + REQUIRE(commands.commands.front().bindings.size() == 10); +} + static void run_qwen_flash_attention_dispatch_checks() { static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); @@ -4907,7 +5369,7 @@ static void run_qwen_flash_attention_dispatch_checks() { { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2); - schedule_common_flash_attention_decode_split_command(ctx, output, 4, 8, 4, 2); + schedule_common_flash_attention_decode_fallback_command(ctx, output, 4, 8, 4, 2); } { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 16, 16, 4, 2); @@ -4942,7 +5404,7 @@ static void run_qwen_flash_attention_dispatch_checks() { } { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 8, 4, 2); - schedule_common_flash_attention_decode_split_command(ctx, output, 1, 8, 4, 2); + schedule_common_flash_attention_decode_fallback_command(ctx, output, 1, 8, 4, 2); } { REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); @@ -4951,8 +5413,16 @@ static void run_qwen_flash_attention_dispatch_checks() { ggml_tensor * output = build_qwen_flash_attention_graph( ctx, 1, 512, 32, 4, GGML_TYPE_F32, GGML_TYPE_F16, true, false, qk_head_size, 1.0f / std::sqrt(static_cast(qk_head_size)), value_head_size); - schedule_common_flash_attention_decode_split_command(ctx, output, 1, 512, 32, 4, qk_head_size, value_head_size, - 1.0f / std::sqrt(static_cast(qk_head_size))); + schedule_common_flash_attention_decode_fallback_command(ctx, output, 1, 512, 32, 4, qk_head_size, + value_head_size, + 1.0f / std::sqrt(static_cast(qk_head_size))); + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + schedule_common_flash_attention_decode_q8_command(ctx, GGML_TYPE_Q4_K); + schedule_common_flash_attention_decode_q8_command(ctx, GGML_TYPE_Q6_K); + schedule_common_flash_attention_decode_unqualified_command(ctx); REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); } { @@ -5449,19 +5919,20 @@ static void run_packed_f16_producer_consumer_checks() { int64_t outputs; bool side_use; bool packed; + bool alternate; bool packed_copy = false; } cases[] = { - { GGML_TYPE_Q6_K, 512, 4096, 2048, false, true }, - { GGML_TYPE_Q6_K, 1024, 4096, 2048, false, true }, - { GGML_TYPE_Q6_K, 256, 4096, 2048, false, false }, - { GGML_TYPE_Q6_K, 512, 4096, 1024, false, false }, - { GGML_TYPE_Q4_K, 512, 4096, 4096, false, true }, - { GGML_TYPE_Q4_K, 512, 8192, 2048, false, false }, - { GGML_TYPE_Q4_K, 512,12288, 4096, false, true }, - { GGML_TYPE_Q4_K, 512,12288, 4096, true, false }, - { GGML_TYPE_Q4_K, 512,32768, 4096, false, false }, - { GGML_TYPE_Q4_K, 1024, 4096, 4096, false, false }, - { GGML_TYPE_Q6_K, 512, 4096, 2048, true, false, true }, + { GGML_TYPE_Q6_K, 512, 4096, 2048, false, true, false }, + { GGML_TYPE_Q6_K, 1024, 4096, 2048, false, true, false }, + { GGML_TYPE_Q6_K, 256, 4096, 2048, false, false, true }, + { GGML_TYPE_Q6_K, 512, 4096, 1024, false, false, true }, + { GGML_TYPE_Q4_K, 512, 4096, 4096, false, true, false }, + { GGML_TYPE_Q4_K, 512, 8192, 2048, false, false, true }, + { GGML_TYPE_Q4_K, 512,12288, 4096, false, true, false }, + { GGML_TYPE_Q4_K, 512,12288, 4096, true, false, true }, + { GGML_TYPE_Q4_K, 512,32768, 4096, false, false, false }, + { GGML_TYPE_Q4_K, 1024, 4096, 4096, false, false, true }, + { GGML_TYPE_Q6_K, 512, 4096, 2048, true, false, true, true }, }; for (const auto & test : cases) { @@ -5490,7 +5961,7 @@ static void run_packed_f16_producer_consumer_checks() { plan.metadata.find_generated_resource(value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); const auto * alternate = plan.metadata.find_alternate_value(value->id, GGML_TYPE_F16, value->byte_count / 2); REQUIRE((generated != nullptr) == (test.packed || test.packed_copy)); - REQUIRE((alternate != nullptr) != test.packed); + REQUIRE((alternate != nullptr) == test.alternate); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"; @@ -5520,6 +5991,7 @@ static void run_packed_f16_producer_consumer_checks() { REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } + ggml_free(ctx); } @@ -5583,6 +6055,122 @@ static void run_generic_swiglu_k16_publish_checks() { ggml_free(ctx); } +static void run_generic_swiglu_k16_publication_edge_checks() { + { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t input_size = 2048; + constexpr int64_t hidden_size = 8192; + constexpr int64_t token_count = 512; + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + ggml_tensor * glu = ggml_swiglu_split(ctx, gate, up); + REQUIRE(glu != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, glu); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); + REQUIRE(glu_value != nullptr); + REQUIRE(plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == + nullptr); + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings.back().value == glu_value->id); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t input_size = 2048; + constexpr int64_t hidden_size = 8192; + constexpr int64_t token_count = 512; + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + ggml_tensor * glu = ggml_swiglu_split(ctx, gate, up); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, glu, 4096, 1024); + ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 4096, 2048); + ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, reshaped); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); + const ggml::hrx::Value * reshaped_value = imported.graph.values().find_tensor(reshaped); + REQUIRE(glu_value != nullptr); + REQUIRE(reshaped_value != nullptr); + REQUIRE(plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == + nullptr); + const auto * packed = + plan.metadata.find_generated_resource(reshaped_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(packed != nullptr); + + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 4); + const auto copy = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major" && + dispatch.bindings.front().value == reshaped_value->id; + }); + REQUIRE(copy != plan.dispatches.end()); + REQUIRE(copy->bindings.back().value == packed->generated_value); + require_compile_parameter(*copy, "ggml.copy_f16_k16_major.input_is_f16", "0"); + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*copy && !dispatch.bindings.empty() && + dispatch.bindings.front().value == packed->generated_value; + }); + REQUIRE(consumer != plan.dispatches.end()); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + static void run_tiled_matmul_alternate_publish_checks() { struct Case { ggml_type consumer_type; @@ -5660,6 +6248,93 @@ static void run_tiled_matmul_alternate_publish_checks() { REQUIRE(command_program_verifies(commands)); ggml_free(ctx); } + + { + ggml_init_params params = {}; + params.mem_size = 12 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); + ggml_tensor * producer_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 2048); + ggml_tensor * projection = ggml_mul_mat(ctx, producer_weight, input); + ggml_tensor * q5_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, 2048, 1024); + ggml_tensor * q6_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 1024); + ggml_tensor * q5_output = ggml_mul_mat(ctx, q5_weight, projection); + ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, projection); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, q5_output); + ggml_build_forward_expand(graph, q6_output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const auto * projection_value = imported.graph.values().find_tensor(projection); + REQUIRE(projection_value != nullptr); + const size_t q8_bytes = static_cast(256) * ggml_row_size(GGML_TYPE_Q8_1, 2048); + const auto * q8 = plan.metadata.find_alternate_value(projection_value->id, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(q8 != nullptr); + const auto publication_count = std::count_if( + plan.metadata.activation_publication_diagnostics().begin(), + plan.metadata.activation_publication_diagnostics().end(), [&](const auto & diagnostic) { + return diagnostic.source_value == projection_value->id; + }); + REQUIRE(publication_count == 1); + const auto publication = std::find_if( + plan.metadata.activation_publication_diagnostics().begin(), + plan.metadata.activation_publication_diagnostics().end(), [&](const auto & diagnostic) { + return diagnostic.source_value == projection_value->id; + }); + REQUIRE(publication != plan.metadata.activation_publication_diagnostics().end()); + REQUIRE(publication->requested_format == "q8-1-x4"); + + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" && + dispatch.bindings[2].value == projection_value->id; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 3); + + const auto q5_consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256" && + dispatch.bindings.front().value == q8->alternate_value; + }); + REQUIRE(q5_consumer != plan.dispatches.end()); + const auto q6_consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"; + }); + REQUIRE(q6_consumer != plan.dispatches.end()); + const auto f16_copy = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f32_f16" && + dispatch.bindings.front().value == projection_value->id; + }); + REQUIRE(f16_copy != plan.dispatches.end()); + REQUIRE(q6_consumer->bindings.front().value == f16_copy->bindings.back().value); + + const size_t q8_quantizer_count = std::count_if( + plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && + dispatch.bindings.front().value == projection_value->id; + }); + REQUIRE(q8_quantizer_count == 1); + + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } } static void run_tiled_pair_matmul_postops_dispatch_checks() { @@ -5926,6 +6601,7 @@ static void run_rmsnorm_binary_k16_output_checks() { bool packed; } cases[] = { { GGML_TYPE_Q4_K, 2048, 512, 1, false, false, false, true, true }, + { GGML_TYPE_Q4_K, 3072, 512, 1, false, false, false, true, true }, { GGML_TYPE_Q4_K, 4096, 512, 1, false, false, false, true, true }, { GGML_TYPE_Q4_K, 5120, 512, 2, false, false, false, true, true }, { GGML_TYPE_Q6_K, 8192, 512, 2, true, false, false, true, true }, @@ -5995,6 +6671,17 @@ static void run_rmsnorm_binary_k16_output_checks() { return &dispatch != &*producer && !dispatch.bindings.empty() && dispatch.bindings.front().value == generated->generated_value; }) == test.projections); + } else if (test.hidden == 7168 && test.tokens == 512) { + const auto * f16 = + plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2); + REQUIRE(f16 != nullptr); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; + }) == 1); + } else if (test.projections == 0) { + REQUIRE(plan.metadata.find_generated_resource( + output->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == nullptr); + REQUIRE(plan.metadata.find_alternate_value(output->id) == nullptr); } if (side != nullptr) { const auto * side_value = imported.graph.values().find_tensor(side); @@ -6086,9 +6773,10 @@ static void run_rmsnorm_binary_f16_output_checks() { const ggml::hrx::CommandPlan & plan = scheduler.plan(); REQUIRE(plan.valid()); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_f16"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; }); REQUIRE(producer != plan.dispatches.end()); + require_compile_parameter(*producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); REQUIRE(producer->bindings.size() == 5); const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); REQUIRE(norm_value != nullptr); @@ -6114,6 +6802,55 @@ static void run_rmsnorm_binary_f16_output_checks() { REQUIRE(command_program_verifies(commands)); ggml_free(ctx); } + + { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t hidden = 2048; + constexpr int64_t tokens = 2; + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hidden, 1, tokens); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden); + ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-5f), norm_weight); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, norm, hidden, tokens); + ggml_tensor * projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, hidden, 6144); + ggml_tensor * output = ggml_silu(ctx, ggml_mul_mat(ctx, projection, reshaped)); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; + }); + REQUIRE(producer != plan.dispatches.end()); + require_compile_parameter(*producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); + REQUIRE(producer->bindings.size() == 5); + const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); + REQUIRE(norm_value != nullptr); + const size_t q8_bytes = qwen_q8_1_x4_size(tokens, hidden); + const auto * q8 = plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(q8 != nullptr); + REQUIRE(producer->bindings[3].value == q8->alternate_value); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; + }) == 0); + REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { + return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; + })); + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } } static void run_rmsnorm_gate_packed_output_checks() { @@ -6138,6 +6875,8 @@ static void run_rmsnorm_gate_packed_output_checks() { { GGML_TYPE_Q4_K, 128, 6144, 512, true, false, 1 }, { GGML_TYPE_Q4_K, 128, 4096, 256, false, false, 0 }, { GGML_TYPE_Q4_K, 256, 256, 512, false, true, 0 }, + { GGML_TYPE_Q5_K, 128, 6144, 19, false, false, 0 }, + { GGML_TYPE_IQ4_XS, 128, 6144, 19, false, false, 0 }, }; for (const auto & test : cases) { @@ -6162,9 +6901,10 @@ static void run_rmsnorm_gate_packed_output_checks() { const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_f16"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; }); REQUIRE(producer != plan.dispatches.end()); + require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "0"); require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.f16_output_row_width", std::to_string(test.packed ? test.channels : 0)); const auto * output = imported.graph.values().find_tensor(gated); @@ -6188,6 +6928,42 @@ static void run_rmsnorm_gate_packed_output_checks() { REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } + + { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 2); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + ggml_tensor * gate = ggml_dup_tensor(ctx, input); + ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-6f), weight); + ggml_tensor * gated = ggml_mul(ctx, norm, ggml_silu(ctx, gate)); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_scale(ctx, gated, 0.5f)); + + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; + }); + REQUIRE(producer != plan.dispatches.end()); + require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "0"); + REQUIRE(producer->bindings.size() == 5); + const auto * output = imported.graph.values().find_tensor(gated); + REQUIRE(output != nullptr); + REQUIRE(plan.metadata.find_alternate_value(output->id) == nullptr); + REQUIRE(plan.metadata.find_generated_resource( + output->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == nullptr); + REQUIRE(std::any_of(plan.transients.begin(), plan.transients.end(), [&](const auto & transient) { + return transient.value == producer->bindings.back().value && + transient.size == output->byte_count / 2; + })); + const auto commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } ggml_free(ctx); } @@ -6248,10 +7024,13 @@ static void run_rmsnorm_gate_q8_output_checks() { const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_q8_1_x4"; + const auto publish_q8 = dispatch.kernel.compile_parameters.find("ggml.rmsnorm_gate_f32.publish_q8"); + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish" && + publish_q8 != dispatch.kernel.compile_parameters.end() && publish_q8->second == "1"; }); REQUIRE((producer != plan.dispatches.end()) == test.q8); if (test.q8) { + require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "1"); const auto * output = imported.graph.values().find_tensor(gated); REQUIRE(output != nullptr); const size_t bytes = static_cast(test.tokens) * ggml_row_size(GGML_TYPE_Q8_1, test.channels); @@ -6275,6 +7054,105 @@ static void run_rmsnorm_gate_q8_output_checks() { ggml_free(ctx); } +static void run_symmetric_i4_consumer_qualification_checks() { + const struct { + bool direct_consumer; + bool reshape_consumer; + bool strided_consumer; + bool q8_publication; + } cases[] = { + { true, false, false, true }, + { false, true, false, true }, + { false, false, true, false }, + { false, false, false, false }, + }; + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t hidden = 2048; + constexpr int64_t tokens = 2; + const int64_t producer_hidden = test.strided_consumer ? hidden + 128 : hidden; + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, producer_hidden, tokens); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, producer_hidden); + ggml_tensor * normalized = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-6f), weight); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_tensor * consumer_input = normalized; + if (test.reshape_consumer) { + consumer_input = ggml_reshape_2d(ctx, normalized, hidden, tokens); + } else if (test.strided_consumer) { + consumer_input = ggml_view_2d( + ctx, normalized, hidden, tokens, static_cast(producer_hidden) * sizeof(float), 0); + } + if (test.direct_consumer || test.reshape_consumer || test.strided_consumer) { + ggml_tensor * i4_projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, hidden, hidden); + ggml_tensor * i4_projected = ggml_mul_mat(ctx, i4_projection, consumer_input); + ggml_build_forward_expand(graph, ggml_scale(ctx, i4_projected, 0.5f)); + if (!test.strided_consumer) { + ggml_tensor * q8_projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden, 512); + ggml_tensor * q8_projected = ggml_mul_mat(ctx, q8_projection, consumer_input); + ggml_build_forward_expand(graph, ggml_scale(ctx, q8_projected, 0.5f)); + } + } else { + ggml_build_forward_expand(graph, ggml_scale(ctx, normalized, 0.5f)); + } + + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); + const ggml::hrx::Value * qualified_input = imported.graph.values().find_tensor(consumer_input); + REQUIRE(normalized_value != nullptr); + REQUIRE(qualified_input != nullptr); + for (const ggml::hrx::GraphNode * consumer : imported.graph.index().consumers(qualified_input->id)) { + REQUIRE(consumer == nullptr || + !ggml::hrx::common_symmetric_i4_lowrow_mul_mat_eligible( + imported.graph, *consumer, *qualified_input)); + } + if (test.strided_consumer) { + ggml_free(ctx); + continue; + } + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const auto symmetric_producer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32"; + }); + REQUIRE(symmetric_producer == plan.dispatches.end()); + const auto q8_producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; + }); + REQUIRE((q8_producer != plan.dispatches.end()) == test.q8_publication); + if (test.q8_publication) { + require_compile_parameter(*q8_producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); + const auto * q8 = plan.metadata.find_alternate_value( + normalized_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(tokens, hidden)); + REQUIRE(q8 != nullptr); + REQUIRE(q8_producer->bindings.size() == 5); + REQUIRE(q8_producer->bindings[3].value == q8->alternate_value); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; + }) == 0); + REQUIRE(std::any_of(std::next(q8_producer), plan.dispatches.end(), [&](const auto & dispatch) { + return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; + })); + } else if (!test.strided_consumer) { + REQUIRE(plan.metadata.alternate_values().empty()); + } + const auto commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } +} + static void run_quantized_conv4_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; @@ -6434,6 +7312,11 @@ static void run_generic_ssm_conv_binary_dispatch_checks() { REQUIRE(plan.dispatches.size() == 1); const auto & dispatch = plan.dispatches.front(); REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_binary_f32"); + REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(dispatch.kernel.integer_parameters.at("n_t") == n_t); + REQUIRE(dispatch.kernel.integer_parameters.at("n_s") == 1); + REQUIRE(dispatch.kernel.compile_parameters.count("llm.ssm_conv.generic.n_t") == 0); + REQUIRE(dispatch.kernel.compile_parameters.count("llm.ssm_conv.generic.n_s") == 0); REQUIRE(dispatch.bindings.size() == 4); require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_op", std::to_string(ggml::hrx::binary_kind_config_value(test.kind))); @@ -6537,6 +7420,7 @@ static void run_gdn_native_projection_pair_dispatch_checks() { {2, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {5, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {16, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, + {17, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 3072, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 4352, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 5120, 16, false, false, false, GGML_TYPE_Q4_K, false}, @@ -6617,6 +7501,16 @@ static void run_gdn_native_projection_pair_dispatch_checks() { REQUIRE(fused->bindings[3].value == imported.graph.values().find_tensor(alpha)->id); REQUIRE(fused->bindings[4].value == imported.graph.values().find_tensor(beta_raw)->id); } + const auto epilogue = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:llm_gated_delta_net_projection_epilogue_f32"; + }); + REQUIRE((epilogue != plan.dispatches.end()) == (test.tokens > 16)); + if (epilogue != plan.dispatches.end()) { + REQUIRE(epilogue->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(epilogue->kernel.integer_parameters.at("element_count") == test.heads * test.tokens); + REQUIRE(epilogue->kernel.compile_parameters.count("llm.gated_delta_net.epilogue_element_count") == 0); + } const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); @@ -6802,7 +7696,7 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { }); REQUIRE((fused != plan.dispatches.end()) == test.fused); REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_f16"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; }) == (test.fused ? 0 : 1)); if (test.fused) { REQUIRE(fused->bindings.size() == 11); @@ -7527,8 +8421,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_input != nullptr); ggml_tensor * iq2_s_output = ggml_mul_mat(ctx, iq2_s_weight, iq2_s_input); REQUIRE(iq2_s_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, - 2048, 640); + schedule_single_matmul_command(ctx, iq2_s_output, + "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", 2, 2048, 640); ggml_tensor * iq2_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ2_S, 2048, 640); ggml_tensor * iq2_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -7536,7 +8430,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_decode_input != nullptr); ggml_tensor * iq2_s_decode_output = ggml_mul_mat(ctx, iq2_s_decode_weight, iq2_s_decode_input); REQUIRE(iq2_s_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + schedule_single_matmul_command(ctx, iq2_s_decode_output, "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32", 1, 2048, 640); ggml_tensor * iq3_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); @@ -7545,8 +8439,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_input != nullptr); ggml_tensor * iq3_s_output = ggml_mul_mat(ctx, iq3_s_weight, iq3_s_input); REQUIRE(iq3_s_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, - 2048, 128); + schedule_single_matmul_command(ctx, iq3_s_output, + "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32", 4, 2048, 128); ggml_tensor * iq3_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); ggml_tensor * iq3_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -7554,7 +8448,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_decode_input != nullptr); ggml_tensor * iq3_s_decode_output = ggml_mul_mat(ctx, iq3_s_decode_weight, iq3_s_decode_input); REQUIRE(iq3_s_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + schedule_single_matmul_command(ctx, iq3_s_decode_output, "loom_libs:ggml_mul_mat_vector_iq3_s_f32_f32", 1, 2048, 128); ggml_tensor * iq4_nl_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); @@ -7827,6 +8721,148 @@ static void run_qwen_matmul_dispatch_checks() { ggml_free(ctx); } +static void run_iq3_xxs_codebook_matmul_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + auto check_plan = [&](int64_t token_count, const char * expected_kernel, size_t expected_binding_count, + const char * weight_format_parameter) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_XXS, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(plan.dispatches[0].kernel.kernel_id) == expected_kernel); + REQUIRE(plan.dispatches[0].bindings.size() == expected_binding_count); + REQUIRE(plan.dispatches[0].bindings[2].length == 1024); + REQUIRE(plan.transients.size() == 1); + REQUIRE(plan.constant_initializations.size() == 1); + REQUIRE(plan.constant_initializations[0].value == plan.dispatches[0].bindings[2].value); + REQUIRE(plan.constant_initializations[0].data.size() == 1024); + require_compile_parameter(plan.dispatches[0], weight_format_parameter, "18"); + }; + + check_plan(1, "loom_libs:ggml_mul_mat_vector_iq3_xxs_f32_f32", 5, "ggml.matmul.vector.weight_format"); + check_plan(64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", 4, + "ggml.mul_mat.weight_format"); + + const char * disable_name = "GGML_HRX_DISABLE_IQ3_XXS_CODEBOOK_MATMUL"; + REQUIRE(setenv(disable_name, "1", 1) == 0); + ggml_tensor * disabled_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_XXS, 2048, 128); + ggml_tensor * disabled_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(disabled_weight != nullptr); + REQUIRE(disabled_input != nullptr); + ggml_tensor * disabled_output = ggml_mul_mat(ctx, disabled_weight, disabled_input); + REQUIRE(disabled_output != nullptr); + ggml_cgraph * disabled_graph = ggml_new_graph(ctx); + REQUIRE(disabled_graph != nullptr); + ggml_build_forward_expand(disabled_graph, disabled_output); + ggml::hrx::GraphImportResult disabled_imported = ggml::hrx::import_ggml_graph(*disabled_graph); + REQUIRE(disabled_imported.valid()); + ggml::hrx::DispatchScheduler disabled_scheduler; + REQUIRE(!disabled_scheduler.schedule_graph(disabled_imported.graph, test_dispatch_target())); + REQUIRE(unsetenv(disable_name) == 0); + + ggml_free(ctx); +} + +static void run_iq2_codebook_resource_reuse_checks() { + for (const ggml_type type : { GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS }) { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * vector_weight = ggml_new_tensor_2d(ctx, type, 2048, 128); + ggml_tensor * vector_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * tiled_weight = ggml_new_tensor_2d(ctx, type, 2048, 128); + ggml_tensor * tiled_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 64); + REQUIRE(vector_weight != nullptr); + REQUIRE(vector_input != nullptr); + REQUIRE(tiled_weight != nullptr); + REQUIRE(tiled_input != nullptr); + ggml_tensor * vector_output = ggml_mul_mat(ctx, vector_weight, vector_input); + ggml_tensor * tiled_output = ggml_mul_mat(ctx, tiled_weight, tiled_input); + REQUIRE(vector_output != nullptr); + REQUIRE(tiled_output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, vector_output); + ggml_build_forward_expand(graph, tiled_output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + REQUIRE(scheduler.plan().constant_initializations.size() == 1); + REQUIRE(scheduler.plan().transients.size() == 1); + + const char * expected_key = + type == GGML_TYPE_IQ2_XXS ? "common.iq2_xxs.grid.v1" : "common.iq2_xs.grid.v1"; + const ggml::hrx::CommandPlanConstantInitialization & initialization = + scheduler.plan().constant_initializations.front(); + REQUIRE(initialization.name == expected_key); + REQUIRE(scheduler.plan().transients.front().name == expected_key); + REQUIRE(scheduler.plan().transients.front().value == initialization.value); + REQUIRE(std::count_if(scheduler.plan().dispatches.begin(), scheduler.plan().dispatches.end(), + [&](const ggml::hrx::Dispatch & dispatch) { + return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), + [&](const ggml::hrx::DispatchBinding & binding) { + return binding.value == initialization.value; + }); + }) == 2); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.constant_initializations.size() == 1); + ggml_free(ctx); + } +} + +static void run_iq4_matmul_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + const auto check = [ctx](ggml_type weight_type, int64_t input_size, int64_t output_size, + int64_t token_count, const char * expected_kernel) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, expected_kernel, token_count, input_size, output_size); + }; + + check(GGML_TYPE_IQ4_NL, 640, 1024, 1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64"); + check(GGML_TYPE_IQ4_NL, 640, 1024, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"); + check(GGML_TYPE_IQ4_XS, 2048, 128, 1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64"); + check(GGML_TYPE_IQ4_XS, 2048, 128, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"); + + ggml_free(ctx); +} + static void run_quantized_value_projection_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 16 * 1024 * 1024; @@ -8417,10 +9453,11 @@ static QwenRoutedGateUpTensors build_qwen_routed_gate_up_graph(ggml_context * ct static void append_qwen_weighted_reduce_tail(ggml_context * ctx, QwenRoutedGateUpTensors & tensors, - bool include_next_rmsnorm = false) { + bool include_next_rmsnorm = false, + ggml_tensor * routed_input = nullptr) { REQUIRE(tensors.output != nullptr); REQUIRE(tensors.route_weights != nullptr); - tensors.weighted = ggml_mul(ctx, tensors.output, tensors.route_weights); + tensors.weighted = ggml_mul(ctx, routed_input != nullptr ? routed_input : tensors.output, tensors.route_weights); REQUIRE(tensors.weighted != nullptr); tensors.route_views.clear(); for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { @@ -8487,6 +9524,7 @@ static bool match_dispatch_at_index(const ggml::hrx::Graph & graph, plan, ggml::hrx::ValueId(static_cast(graph.values().size() + plan.transients.size() + plan.completion_counter_requests.size())), + &test_dispatch_registry(), }; return test_dispatch_registry().match(context, match); } @@ -9219,9 +10257,9 @@ static void run_qwen_routed_gate_up_dispatch_checks() { { constexpr int64_t token_count = 4; - const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors); ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); - REQUIRE(imported.graph.nodes().size() == 14); const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); const ggml::hrx::Value * route_weights_value = imported.graph.values().find_tensor(tensors.route_weights); @@ -9230,7 +10268,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(route_weights_value != nullptr); REQUIRE(glu_value != nullptr); REQUIRE(route_ids_value->kind == ggml::hrx::ValueKind::Transient); - REQUIRE(route_weights_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(route_weights_value->kind == ggml::hrx::ValueKind::Transient); REQUIRE(glu_value->kind == ggml::hrx::ValueKind::Transient); std::vector covered_nodes(imported.graph.nodes().size(), false); @@ -9320,7 +10358,8 @@ static void run_qwen_routed_gate_up_dispatch_checks() { { constexpr int64_t token_count = 1; - const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors); 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); @@ -9350,12 +10389,13 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(first_route_ids != nullptr); REQUIRE(first_route_weights != nullptr); - const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, first_route_weights); ggml_build_forward_expand(graph, tensors.route_weights); - ggml_build_forward_expand(graph, tensors.output); + ggml_build_forward_expand(graph, tensors.residual); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); @@ -9425,17 +10465,20 @@ static void run_qwen_routed_gate_up_dispatch_checks() { std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); - append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, - "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); + REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(token_count)); + 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); REQUIRE(plan.dispatches.size() == 5); - REQUIRE(plan.transients.size() == 4); + REQUIRE(plan.transients.size() == 2); + REQUIRE(plan.metadata.alternate_values().empty()); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 5); REQUIRE(command_program_verifies(commands)); - REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, plan.transients.back().value) != nullptr); } { @@ -9489,6 +10532,66 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(command_program_verifies(commands)); } + { + constexpr int64_t token_count = 4; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + ggml_tensor * glu_reshape = ggml_reshape_4d(ctx, tensors.glu, tensors.glu->ne[0], tensors.glu->ne[1], + tensors.glu->ne[2], tensors.glu->ne[3]); + REQUIRE(glu_reshape != nullptr); + tensors.output = ggml_mul_mat_id(ctx, tensors.down_weight, glu_reshape, tensors.route_ids); + REQUIRE(tensors.output != nullptr); + ggml_tensor * routed_output_view = + ggml_view_4d(ctx, tensors.output, tensors.output->ne[0], tensors.output->ne[1], tensors.output->ne[2], + tensors.output->ne[3], tensors.output->nb[1], tensors.output->nb[2], tensors.output->nb[3], 0); + REQUIRE(routed_output_view != nullptr); + append_qwen_weighted_reduce_tail(ctx, tensors, false, routed_output_view); + + ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + const ggml::hrx::Value * glu_reshape_value = imported.graph.values().find_tensor(glu_reshape); + const ggml::hrx::Value * routed_output_view_value = imported.graph.values().find_tensor(routed_output_view); + REQUIRE(glu_reshape_value != nullptr); + REQUIRE(routed_output_view_value != nullptr); + const ggml::hrx::GraphNode * glu_reshape_node = imported.graph.index().producer(glu_reshape_value->id); + const ggml::hrx::GraphNode * routed_output_view_node = + imported.graph.index().producer(routed_output_view_value->id); + REQUIRE(glu_reshape_node != nullptr); + REQUIRE(routed_output_view_node != nullptr); + REQUIRE(glu_reshape_node->op == GGML_OP_RESHAPE); + REQUIRE(routed_output_view_node->op == GGML_OP_VIEW); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(tensors.glu); + REQUIRE(glu_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * gate_up_alternate = ggml::hrx::find_alternate_value( + imported.graph, plan, glu_value->id, GGML_TYPE_F16, qwen_routed_gate_up_f16_output_size(token_count)); + REQUIRE(gate_up_alternate != nullptr); + const ggml::hrx::ValueId gate_up_alternate_value = gate_up_alternate->alternate_value; + + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); + REQUIRE(plan.dispatches.back().bindings[0].value == gate_up_alternate_value); + const ggml::hrx::Value * routed_output_value = imported.graph.values().find_tensor(tensors.output); + REQUIRE(routed_output_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * routed_down_alternate = ggml::hrx::find_alternate_value( + imported.graph, plan, routed_output_value->id, GGML_TYPE_F16, + qwen_routed_down_f16_output_size(token_count)); + REQUIRE(routed_down_alternate != nullptr); + const ggml::hrx::ValueId routed_down_alternate_value = routed_down_alternate->alternate_value; + + append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); + REQUIRE(plan.dispatches.back().bindings[1].value == routed_down_alternate_value); + for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != "loom_libs:ggml_copy_f32_f16"); + } + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + { constexpr int64_t token_count = 4; QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); @@ -9563,8 +10666,9 @@ static void run_qwen_routed_gate_up_dispatch_checks() { { constexpr int64_t token_count = 4; - const QwenRoutedGateUpTensors tensors = + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_Q4_K); + append_qwen_weighted_reduce_tail(ctx, tensors); 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); @@ -9678,6 +10782,89 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml_free(ctx); } +static void run_qwen_weighted_reduce_next_rmsnorm_q8_publication_checks() { + enum class ConsumerCase { + Direct, + Alias, + UnsupportedWeight, + WrongOperand, + IncompatibleGeometry, + }; + const struct { + ConsumerCase consumer; + bool publishes_q8; + } cases[] = { + { ConsumerCase::Direct, true }, + { ConsumerCase::Alias, true }, + { ConsumerCase::UnsupportedWeight, false }, + { ConsumerCase::WrongOperand, false }, + { ConsumerCase::IncompatibleGeometry, false }, + }; + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t token_count = 2; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); + append_qwen_weighted_reduce_tail(ctx, tensors, true); + ggml_tensor * consumer_input = tensors.next_output; + if (test.consumer == ConsumerCase::Alias) { + consumer_input = ggml_reshape_2d(ctx, tensors.next_output, kQwenMoeHiddenSize, token_count); + REQUIRE(consumer_input != nullptr); + } + + if (test.consumer == ConsumerCase::WrongOperand) { + ggml_tensor * activation = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, 64); + REQUIRE(activation != nullptr); + tensors.hidden_use = ggml_mul_mat(ctx, tensors.next_output, activation); + } else { + const ggml_type weight_type = + test.consumer == ConsumerCase::UnsupportedWeight ? GGML_TYPE_BF16 : GGML_TYPE_Q4_K; + const int64_t output_size = test.consumer == ConsumerCase::IncompatibleGeometry ? 63 : 512; + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, kQwenMoeHiddenSize, output_size); + REQUIRE(weight != nullptr); + tensors.hidden_use = ggml_mul_mat(ctx, weight, consumer_input); + } + REQUIRE(tensors.hidden_use != nullptr); + + 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); + append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); + + const size_t weighted_index = producer_index_for_tensor(imported.graph, tensors.weighted); + ggml::hrx::DispatchMatch weighted_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, weighted_index, weighted_match)); + const char * expected_kernel = + test.publishes_q8 ? + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8" : + "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"; + REQUIRE(kernel_name_for_id(weighted_match.dispatches.back().kernel.kernel_id) == expected_kernel); + REQUIRE(weighted_match.dispatches.back().bindings.size() == (test.publishes_q8 ? 6 : 5)); + + const ggml::hrx::Value * next_output_value = imported.graph.values().find_tensor(tensors.next_output); + REQUIRE(next_output_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * q8 = weighted_match.metadata.find_alternate_value( + next_output_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(token_count, kQwenMoeHiddenSize)); + REQUIRE((q8 != nullptr) == test.publishes_q8); + if (test.publishes_q8) { + REQUIRE(weighted_match.metadata.activation_publication_diagnostics().size() == 1); + const auto & diagnostic = weighted_match.metadata.activation_publication_diagnostics().front(); + REQUIRE(diagnostic.source_value == next_output_value->id); + REQUIRE(diagnostic.alternate_value == q8->alternate_value); + REQUIRE(diagnostic.publication_stage == "published"); + REQUIRE(diagnostic.requested_format == "q8-1-x4"); + } + ggml_free(ctx); + } +} + static void run_qwen_router_top8_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 2 * 1024 * 1024; @@ -11218,6 +12405,214 @@ static void write_transient_i32_value(ggml_backend_hrx_context * con hrx_synchronous_h2d(context->device->device, &value, arena.buffer, allocation.arena_offset, sizeof(value))); } +static std::vector read_transient_bytes(ggml_backend_hrx_context * context, + ggml::hrx::TransientArenaAllocationRef arena, + const ggml::hrx::TransientAllocation & allocation) { + REQUIRE(context != nullptr); + REQUIRE(arena.buffer != nullptr); + std::vector data(allocation.size); + require_hrx_status(hrx_synchronous_d2h(context->device->device, arena.buffer, allocation.arena_offset, data.data(), + allocation.size)); + return data; +} + +static void run_binary_q8_publication_numerics() { + const struct { + int64_t hidden; + int64_t tokens; + bool broadcast; + int op; + } cases[] = { + { 2048, 3, false, 0 }, + { 2176, 4, false, 2 }, + { 2048, 4, true, 0 }, + { 2176, 5, true, 2 }, + }; + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto * backend_context = static_cast(backend->context); + REQUIRE(backend_context != nullptr); + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + const char * target = backend_context->device->architecture.c_str(); + + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + const int64_t element_count = test.hidden * test.tokens; + const int64_t rhs_count = test.broadcast ? test.hidden : element_count; + ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); + ggml_tensor * rhs = test.broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden) : + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); + REQUIRE(lhs != nullptr && rhs != nullptr); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector lhs_data(element_count); + std::vector rhs_data(rhs_count); + std::vector expected(element_count); + for (int64_t i = 0; i < element_count; ++i) { + lhs_data[i] = static_cast((i * 17 + 3) % 251 - 125) / 31.0f; + } + for (int64_t i = 0; i < rhs_count; ++i) { + rhs_data[i] = static_cast((i * 29 + 11) % 239 - 119) / 37.0f; + } + for (int64_t token = 0; token < test.tokens; ++token) { + for (int64_t channel = 0; channel < test.hidden; ++channel) { + const int64_t i = token * test.hidden + channel; + const float rhs_v = rhs_data[test.broadcast ? channel : i]; + expected[i] = test.op == 0 ? lhs_data[i] + rhs_v : lhs_data[i] * rhs_v; + } + } + ggml_backend_tensor_set(lhs, lhs_data.data(), 0, lhs_data.size() * sizeof(float)); + ggml_backend_tensor_set(rhs, rhs_data.data(), 0, rhs_data.size() * sizeof(float)); + ggml_backend_synchronize(backend); + + ggml::hrx::Graph graph; + const ggml::hrx::ValueId lhs_value = + graph.values().get_or_add_tensor_value(lhs, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId rhs_value = + graph.values().get_or_add_tensor_value(rhs, ggml::hrx::ValueKind::External); + ggml::hrx::ValueBufferBinding lhs_binding; + ggml::hrx::ValueBufferBinding rhs_binding; + REQUIRE(ggml_backend_hrx_resolve_value_buffer(lhs, lhs_binding)); + REQUIRE(ggml_backend_hrx_resolve_value_buffer(rhs, rhs_binding)); + REQUIRE(graph.values().bind_buffer(lhs_value, lhs_binding)); + REQUIRE(graph.values().bind_buffer(rhs_value, rhs_binding)); + + const ggml::hrx::ValueId reference_f32(static_cast(graph.values().size())); + const ggml::hrx::ValueId reference_q8(reference_f32.value + 1); + const ggml::hrx::ValueId published_f32(reference_f32.value + 2); + const ggml::hrx::ValueId published_q8(reference_f32.value + 3); + const size_t f32_bytes = static_cast(element_count) * sizeof(float); + const size_t q8_bytes = qwen_q8_1_x4_size(test.tokens, test.hidden); + + ggml::hrx::CommandPlan plan; + plan.transients.push_back({ reference_f32, "test.binary.reference_f32", f32_bytes, 256 }); + plan.transients.push_back({ reference_q8, "test.binary.reference_q8", q8_bytes, 256 }); + plan.transients.push_back({ published_f32, "test.binary.published_f32", f32_bytes, 256 }); + plan.transients.push_back({ published_q8, "test.binary.published_q8", q8_bytes, 256 }); + + ggml::hrx::Dispatch legacy; + legacy.kernel = ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref( + "loom_libs", test.broadcast ? "ggml_binary_bc_f32" : "ggml_binary_f32")); + legacy.kernel.integer_parameters.emplace("element_count", element_count); + legacy.kernel.compile_parameters.emplace( + test.broadcast ? "ggml.binary_bc_f32.op" : "ggml.binary_f32.op", std::to_string(test.op)); + if (test.broadcast) { + legacy.kernel.integer_parameters.emplace("ne0", test.hidden); + legacy.kernel.integer_parameters.emplace("ne1", test.tokens); + legacy.kernel.integer_parameters.emplace("ne2", 1); + legacy.kernel.integer_parameters.emplace("ne3", 1); + legacy.kernel.integer_parameters.emplace("src0_element_count", element_count); + legacy.kernel.integer_parameters.emplace("src1_element_count", rhs_count); + } else { + 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_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_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)); + legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", std::to_string(element_count)); + legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", std::to_string(element_count)); + } + if (test.broadcast) { + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + legacy.kernel.compile_parameters.emplace("ggml.binary_bc_f32.src0_broadcast_dim" + + std::to_string(dim), + "0"); + legacy.kernel.compile_parameters.emplace("ggml.binary_bc_f32.src1_broadcast_dim" + + std::to_string(dim), + dim == 1 ? "1" : "0"); + } + } + legacy.bindings.push_back({ lhs_value, 0, ggml_nbytes(lhs) }); + legacy.bindings.push_back({ rhs_value, 0, ggml_nbytes(rhs) }); + legacy.bindings.push_back({ reference_f32, 0, f32_bytes }); + plan.dispatches.push_back(std::move(legacy)); + + ggml::hrx::Dispatch quantize; + quantize.kernel = ggml::hrx::make_kernel_specialization( + ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_quantize_q8_1_x4_f32")); + quantize.kernel.integer_parameters.emplace("token_count", test.tokens); + quantize.kernel.integer_parameters.emplace("input_size", test.hidden); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + std::to_string(element_count / 128)); + quantize.bindings.push_back({ reference_f32, 0, f32_bytes }); + quantize.bindings.push_back({ reference_q8, 0, q8_bytes }); + plan.dispatches.push_back(std::move(quantize)); + + ggml::hrx::Dispatch publish; + publish.kernel = ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref( + "loom_libs", test.broadcast ? "ggml_binary_bc_f32_publish_q8_1_x4" : + "ggml_binary_f32_publish_q8_1_x4")); + publish.kernel.integer_parameters.emplace("token_count", test.tokens); + if (test.broadcast) { + publish.kernel.integer_parameters.emplace("hidden_size", test.hidden); + publish.kernel.integer_parameters.emplace("src0_element_count", element_count); + publish.kernel.integer_parameters.emplace("src1_element_count", rhs_count); + publish.kernel.compile_parameters = plan.dispatches.front().kernel.compile_parameters; + } else { + publish.kernel.compile_parameters = plan.dispatches.front().kernel.compile_parameters; + } + publish.bindings.push_back({ lhs_value, 0, ggml_nbytes(lhs) }); + publish.bindings.push_back({ rhs_value, 0, ggml_nbytes(rhs) }); + publish.bindings.push_back({ published_f32, 0, f32_bytes }); + publish.bindings.push_back({ published_q8, 0, q8_bytes }); + plan.dispatches.push_back(std::move(publish)); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(graph, plan, corpus, target); + REQUIRE(commands.valid()); + REQUIRE(ggml::hrx::verify_command_program(commands, corpus, target).valid()); + const ggml::hrx::CommandProgramBindings bindings = + ggml::hrx::CommandProgramBindings::from_value_map(graph.values()); + REQUIRE(bindings.valid()); + const ggml::hrx::CommandProgramExecutionContext execution_context = { + backend_context->device->device, + backend_context->stream, + target, + &corpus, + &backend_context->kernel_executables, + &backend_context->transient_arena, + &backend_context->host_transfers, + &backend_context->host_weights, + }; + REQUIRE(ggml::hrx::execute_command_program(execution_context, commands, bindings)); + ggml_backend_synchronize(backend); + + const auto arena = backend_context->transient_arena.current_allocation(); + const auto * reference_f32_allocation = + ggml::hrx::find_transient_allocation(commands.transients, reference_f32); + const auto * reference_q8_allocation = + ggml::hrx::find_transient_allocation(commands.transients, reference_q8); + const auto * published_f32_allocation = + ggml::hrx::find_transient_allocation(commands.transients, published_f32); + const auto * published_q8_allocation = + ggml::hrx::find_transient_allocation(commands.transients, published_q8); + REQUIRE(reference_f32_allocation != nullptr && reference_q8_allocation != nullptr); + REQUIRE(published_f32_allocation != nullptr && published_q8_allocation != nullptr); + const auto reference_f32_bytes = read_transient_bytes(backend_context, arena, *reference_f32_allocation); + const auto reference_q8_bytes = read_transient_bytes(backend_context, arena, *reference_q8_allocation); + const auto published_f32_bytes = read_transient_bytes(backend_context, arena, *published_f32_allocation); + const auto published_q8_bytes = read_transient_bytes(backend_context, arena, *published_q8_allocation); + REQUIRE(reference_f32_bytes == published_f32_bytes); + REQUIRE(reference_q8_bytes == published_q8_bytes); + REQUIRE(reference_f32_bytes.size() == expected.size() * sizeof(float)); + REQUIRE(std::memcmp(reference_f32_bytes.data(), expected.data(), reference_f32_bytes.size()) == 0); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + } + ggml_backend_free(backend); +} + static void run_qwen_expert_table_partition_prefill_512_execution() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -12292,7 +13687,7 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { REQUIRE(std::none_of(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { const auto name = kernel_name_for_id(dispatch.kernel.kernel_id); return name == "loom_libs:ggml_get_rows_f32" || name == "loom_libs:ggml_copy_f32" || - name == "loom_libs:ggml_rmsnorm_gate_f32_q8_1_x4"; + name == "loom_libs:ggml_rmsnorm_gate_f32_publish"; })); } const auto commands = @@ -12317,6 +13712,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("rmsnorm_two_binary_dispatch", [] { run_rmsnorm_two_binary_dispatch_checks(); }); suite.host_case("cont_f32_dispatch", [] { run_cont_f32_dispatch_checks(); }); suite.host_case("binary_f32_broadcast_dispatch", [] { run_binary_f32_broadcast_dispatch_checks(); }); + suite.host_case("binary_q8_publication", [] { run_binary_q8_publication_checks(); }); suite.host_case("graph_snapshot_diagnostics", [] { run_graph_snapshot_diagnostics_checks(); }); suite.host_case("unmatched_graph_diagnostics", [] { run_unmatched_graph_diagnostics_checks(); }); suite.host_case("completion_counter_plan", [] { run_completion_counter_plan_checks(); }); @@ -12331,8 +13727,13 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("tiled_pair_matmul_postops_dispatch", [] { run_tiled_pair_matmul_postops_dispatch_checks(); }); suite.host_case("qwen_attention_postprocess_dispatch", [] { run_qwen_attention_postprocess_dispatch_checks(); }); suite.host_case("qwen_matmul_dispatch", [] { run_qwen_matmul_dispatch_checks(); }); + suite.host_case("iq1_codebook_matmul_dispatch", [] { run_iq1_codebook_matmul_dispatch_checks(); }); + suite.host_case("iq3_xxs_codebook_matmul_dispatch", [] { run_iq3_xxs_codebook_matmul_dispatch_checks(); }); + suite.host_case("iq2_codebook_resource_reuse", [] { run_iq2_codebook_resource_reuse_checks(); }); + suite.host_case("iq4_matmul_dispatch", [] { run_iq4_matmul_dispatch_checks(); }); suite.host_case("packed_f16_producer_consumer", [] { run_packed_f16_producer_consumer_checks(); }); suite.host_case("generic_swiglu_k16_publish", [] { run_generic_swiglu_k16_publish_checks(); }); + suite.host_case("generic_swiglu_k16_publication_edges", [] { run_generic_swiglu_k16_publication_edge_checks(); }); suite.host_case("tiled_matmul_alternate_publish", [] { run_tiled_matmul_alternate_publish_checks(); }); suite.host_case("k16_major_preparation_reuse", [] { run_k16_major_preparation_reuse_checks(); }); suite.host_case("rmsnorm_binary_k16_output", [] { run_rmsnorm_binary_k16_output_checks(); }); @@ -12340,6 +13741,8 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("swiglu_q8_output", [] { run_swiglu_q8_output_checks(); }); suite.host_case("rmsnorm_gate_packed_output", [] { run_rmsnorm_gate_packed_output_checks(); }); suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); + suite.host_case("symmetric_i4_consumer_qualification", + [] { run_symmetric_i4_consumer_qualification_checks(); }); suite.host_case("lowtoken_residual_dispatch", [] { run_lowtoken_residual_dispatch_checks(); }); suite.host_case("lowtoken_bias_dispatch", [] { run_lowtoken_bias_dispatch_checks(); }); suite.host_case("generic_ssm_conv_binary_dispatch", [] { run_generic_ssm_conv_binary_dispatch_checks(); }); @@ -12356,6 +13759,10 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); }); suite.host_case("qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); }); suite.host_case("qwen_terminal_q6k_vector_command", [] { schedule_qwen_terminal_q6k_vector_command(); }); + suite.host_case("qwen_decode_rmsnorm_publication", [] { run_qwen_decode_rmsnorm_publication_checks(); }); + suite.host_case("qwen_decode_publication_consumer_qualification", [] { + run_qwen_decode_publication_consumer_qualification_checks(); + }); suite.host_case("get_rows_q8_1_alternate_command.q4_k", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); }); suite.host_case("get_rows_q8_1_alternate_command.bf16", @@ -12363,6 +13770,9 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { suite.host_case("qwen_router_top8_dispatch", [] { run_qwen_router_top8_dispatch_checks(); }); suite.host_case("common_mul_mat_id_swiglu_dispatch", [] { run_common_mul_mat_id_swiglu_dispatch_checks(); }); suite.host_case("qwen_routed_gate_up_dispatch", [] { run_qwen_routed_gate_up_dispatch_checks(); }); + suite.host_case("qwen_weighted_reduce_next_rmsnorm_q8_publication", [] { + run_qwen_weighted_reduce_next_rmsnorm_q8_publication_checks(); + }); suite.host_case("common_mul_mat_id_postops_dispatch", [] { run_common_mul_mat_id_postops_dispatch_checks(); }); suite.host_case("alias_value_import", [] { run_alias_value_import_checks(); }); suite.host_case("multi_dispatch", [] { run_multi_dispatch_checks(); }); @@ -12391,6 +13801,7 @@ static void register_hrx_backend_host_cases(test_runner::Suite & suite) { static void register_hrx_backend_device_cases(test_runner::Suite & suite) { suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); + suite.device_case("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(); }); suite.device_case("add_f32", [] { run_add_f32(); }); diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index ed3391a19e47..ee5272eaef95 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -401,25 +401,29 @@ static void run_cache_materialize_case(HrxTestDevice & devic ggml::hrx::Dispatch dispatch = make_binary_add_dispatch(binary, 64); std::vector constants; + ggml_hrx_loom_jit_launch_config launch; ggml::hrx::KernelExecutableCache cache(mode); - ggml::hrx::KernelExecutableRef ref = cache.get_or_compile(context, binary, dispatch, constants); + ggml::hrx::KernelExecutableRef ref = cache.get_or_compile(context, binary, dispatch); REQUIRE(ref.valid()); - std::shared_ptr executable = cache.materialize(context, ref, constants); + std::shared_ptr executable = + cache.materialize(context, ref, dispatch.kernel, constants, launch); REQUIRE(executable != nullptr); REQUIRE(executable->executable != nullptr); REQUIRE(executable->export_info.binding_count == dispatch.bindings.size()); REQUIRE(executable->export_info.constant_byte_length == constants.size()); - REQUIRE(executable->launch.workgroup_count[0] > 0); - REQUIRE(executable->launch.workgroup_size[0] > 0); + REQUIRE(launch.workgroup_count[0] > 0); + REQUIRE(launch.workgroup_size[0] > 0); - std::shared_ptr loaded_hit = cache.materialize(context, ref, constants); + std::shared_ptr loaded_hit = + cache.materialize(context, ref, dispatch.kernel, constants, launch); REQUIRE(loaded_hit == executable); std::vector constants_for_prepare; + ggml_hrx_loom_jit_launch_config launch_for_prepare; std::shared_ptr prepared = - cache.prepare(context, binary, dispatch, constants_for_prepare); + cache.prepare(context, binary, dispatch, constants_for_prepare, launch_for_prepare); REQUIRE(prepared == executable); const char * mode_name = mode == ggml::hrx::LoomJitMode::Async ? "async" : "sync"; @@ -427,6 +431,96 @@ static void run_cache_materialize_case(HrxTestDevice & devic device.architecture.c_str()); } +static void run_dynamic_cache_materialize_case(HrxTestDevice & device, + ggml::hrx::LoomJitMode mode, + const ggml::hrx::KernelDefinition & binary) { + ggml::hrx::KernelExecutablePrepareContext context = {}; + context.device = device.device; + context.target = device.architecture.c_str(); + + ggml::hrx::Dispatch first = make_binary_add_dispatch(binary, 64); + ggml::hrx::Dispatch second = make_binary_add_dispatch(binary, 257); + first.kernel.compile_parameters = binary_f32_exact_config("0", 512); + second.kernel.compile_parameters = first.kernel.compile_parameters; + first.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; + second.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; + + ggml::hrx::KernelExecutableCache cache(mode); + ggml::hrx::KernelExecutableRef first_ref = cache.get_or_compile(context, binary, first); + ggml::hrx::KernelExecutableRef second_ref = cache.get_or_compile(context, binary, second); + REQUIRE(first_ref.valid()); + REQUIRE(second_ref.valid()); + REQUIRE(first_ref.entry == second_ref.entry); + + std::vector first_constants; + std::vector second_constants; + ggml_hrx_loom_jit_launch_config first_launch; + ggml_hrx_loom_jit_launch_config second_launch; + std::shared_ptr first_executable = + cache.materialize(context, first_ref, first.kernel, first_constants, first_launch); + std::shared_ptr second_executable = + cache.materialize(context, second_ref, second.kernel, second_constants, second_launch); + REQUIRE(first_executable != nullptr); + REQUIRE(second_executable == first_executable); + REQUIRE(first_launch.workgroup_count[0] == 1); + REQUIRE(second_launch.workgroup_count[0] == 2); + REQUIRE(first_constants.size() == second_constants.size()); + REQUIRE(first_constants != second_constants); + + ggml::hrx::Dispatch exact_second = second; + exact_second.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Exact; + ggml::hrx::KernelExecutableRef exact_ref = cache.get_or_compile(context, binary, exact_second); + REQUIRE(exact_ref.valid()); + REQUIRE(exact_ref.entry != first_ref.entry); +} + +static void run_dynamic_export_reuse_case( + HrxTestDevice & device, + const char * kernel_name, + const std::map & first_workload, + const std::map & second_workload, + const std::map & compile_config = {}) { + const ggml::hrx::KernelDefinition & definition = find_kernel(kernel_name); + ggml::hrx::KernelExecutablePrepareContext context = {}; + context.device = device.device; + context.target = device.architecture.c_str(); + + auto make_dispatch = [&](const std::map & workload) { + ggml::hrx::Dispatch dispatch; + dispatch.kernel.kernel_id = definition.id; + dispatch.kernel.integer_parameters = workload; + dispatch.kernel.compile_parameters = compile_config; + dispatch.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; + dispatch.bindings.resize(definition.bindings.size()); + return dispatch; + }; + + ggml::hrx::Dispatch first = make_dispatch(first_workload); + ggml::hrx::Dispatch second = make_dispatch(second_workload); + ggml::hrx::KernelExecutableCache cache(ggml::hrx::LoomJitMode::Sync); + ggml::hrx::KernelExecutableRef first_ref = cache.get_or_compile(context, definition, first); + ggml::hrx::KernelExecutableRef second_ref = cache.get_or_compile(context, definition, second); + REQUIRE(first_ref.valid()); + REQUIRE(second_ref.valid()); + REQUIRE(first_ref.entry == second_ref.entry); + + std::vector first_constants; + std::vector second_constants; + ggml_hrx_loom_jit_launch_config first_launch; + ggml_hrx_loom_jit_launch_config second_launch; + std::shared_ptr first_executable = + cache.materialize(context, first_ref, first.kernel, first_constants, first_launch); + std::shared_ptr second_executable = + cache.materialize(context, second_ref, second.kernel, second_constants, second_launch); + REQUIRE(first_executable != nullptr); + REQUIRE(second_executable == first_executable); + REQUIRE(first_launch.workgroup_count[0] > 0); + REQUIRE(second_launch.workgroup_count[0] > 0); + REQUIRE(first_constants.size() == second_constants.size()); + REQUIRE(first_constants != second_constants); + std::printf("dynamic executable reuse materialized %s for two workloads\n", kernel_name); +} + static void run_targeted_export_materialize_case(HrxTestDevice & device) { const ggml::hrx::KernelDefinition * definition = find_targeted_kernel("ggml_linear_q6k_q8_1_x4", device.architecture.c_str()); @@ -449,10 +543,12 @@ static void run_targeted_export_materialize_case(HrxTestDevice & device) { ggml::hrx::KernelExecutableCache cache(ggml::hrx::LoomJitMode::Sync); std::vector constants; - std::shared_ptr executable = cache.prepare(context, *definition, dispatch, constants); + ggml_hrx_loom_jit_launch_config launch; + std::shared_ptr executable = + cache.prepare(context, *definition, dispatch, constants, launch); REQUIRE(executable != nullptr); REQUIRE(executable->executable != nullptr); - REQUIRE(executable->launch.workgroup_count[0] == 151936); + REQUIRE(launch.workgroup_count[0] == 151936); std::printf("KernelExecutableCache materialized targeted %s as export %s for %s\n", definition->symbol, definition->name, device.architecture.c_str()); } @@ -465,7 +561,7 @@ static void run_loom_jit_compile_checks() { const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); const ggml::hrx::KernelDefinition & binary_bc = find_kernel("ggml_binary_bc_f32"); const ggml::hrx::KernelDefinition & gather_add = find_kernel("ggml_gather_add_f32"); - const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("qwen3_moe_rmsnorm_f32"); + const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("ggml_rmsnorm_f32"); const ggml::hrx::KernelDefinition & router_top8 = find_kernel("qwen3_moe_router_top8_f32"); const ggml::hrx::KernelDefinition & expert_table = find_kernel("ggml_moe_build_expert_table"); const ggml::hrx::KernelDefinition & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); @@ -522,10 +618,9 @@ static void run_loom_jit_compile_checks() { { "token_count", 1 } }, { - { "qwen3_moe.model.hidden_size", "2048" }, - { "qwen3_moe.model.rms_epsilon", "0.000001" }, - { "qwen3_moe.workload.token_capacity", "1" }, - { "ggml.quantize_q8_1_x4.group_capacity", "256" }, + { "ggml.rmsnorm_f32.hidden_size", "2048" }, + { "ggml.rmsnorm_f32.input_stride", "2048" }, + { "ggml.rmsnorm_f32.rms_epsilon", "0.000001" }, })); refs.push_back(compile_kernel(*async_jit, "async-router-top8-1", router_top8, { @@ -729,10 +824,10 @@ static void run_loom_jit_compile_checks() { const auto & transpose_f16 = find_kernel("ggml_copy_transpose_f16"); auto transpose_ref = compile_kernel(*sync_jit, "transpose-f16-8192x1024", transpose_f16, { + { "row_count", 8192 }, + { "column_count", 1024 }, }, { - { "ggml.copy_transpose_f16.row_count", "8192" }, - { "ggml.copy_transpose_f16.column_count", "1024" }, }); require_compiled_kernel(transpose_ref); auto transposed_attention_config = flash_attention_config("24", "4", "256", "0.0625"); @@ -1070,6 +1165,19 @@ static void run_moe_routing_compile_checks() { const auto & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); require_compiled_kernel(compile_kernel(*jit, "moe-expert-partition-table", partition_table, partition_table_workload, moe_routing_config())); + + const auto & weighted_reduce_publish_q8 = + find_kernel("qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8"); + REQUIRE(weighted_reduce_publish_q8.bindings.size() == 6); + require_compiled_kernel(compile_kernel(*jit, "moe-weighted-reduce-next-rmsnorm-publish-q8", + weighted_reduce_publish_q8, token_count_workload(2), { + { "qwen3_moe.model.hidden_size", "2048" }, + { "qwen3_moe.model.rms_epsilon", "0.000001" }, + { "qwen3_moe.routed_down.output_size", "2048" }, + { "qwen3_moe.routed_down.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "2" }, + { "ggml.quantize_q8_1_x4.group_capacity", "32" }, + })); } static void run_prefill_fusion_compile_checks() { @@ -1096,12 +1204,13 @@ static void run_prefill_fusion_compile_checks() { { "ggml.rmsnorm_binary_f32.op", "2" }, })); - const auto & rmsnorm_q8_f16 = find_kernel("ggml_rmsnorm_binary_q8_1_x4_f16"); + const auto & rmsnorm_q8_f16 = find_kernel("ggml_rmsnorm_binary_q8_1_x4_publish"); require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-q8-f16", rmsnorm_q8_f16, token_count_workload(512), { { "ggml.rmsnorm_binary_q8_1_x4.hidden_size", "2048" }, { "ggml.rmsnorm_binary_q8_1_x4.rms_epsilon", "0.00001" }, { "ggml.rmsnorm_binary_q8_1_x4.op", "2" }, + { "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1" }, })); const auto & flash_f16 = find_kernel("ggml_flash_attention_f32_f16_wmma_publish_f16"); @@ -1257,12 +1366,12 @@ static void run_binary_fusion_compile_checks() { const auto & ssm = find_kernel("llm_ssm_conv_binary_f32"); require_compiled_kernel(compile_kernel(*jit, "binary-fusion-ssm-rhs-div", ssm, { + { "n_t", 8 }, + { "n_s", 1 }, }, { { "llm.ssm_conv.generic.d_conv", "4" }, { "llm.ssm_conv.generic.d_inner", "64" }, - { "llm.ssm_conv.generic.n_t", "8" }, - { "llm.ssm_conv.generic.n_s", "1" }, { "llm.ssm_conv.generic.unary_op", "23" }, { "llm.ssm_conv.generic.binary_op", "3" }, { "llm.ssm_conv.generic.binary_lhs", "0" }, @@ -1324,9 +1433,111 @@ static void run_loom_jit_materialize_checks() { REQUIRE(device.open()); run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); + run_dynamic_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); + run_dynamic_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); + run_dynamic_export_reuse_case(device, "ggml_unary_f32", { { "element_count", 64 } }, + { { "element_count", 257 } }, { { "ggml.unary_f32.op", "4" } }); + run_dynamic_export_reuse_case(device, "ggml_copy_f32", { { "element_count", 64 } }, + { { "element_count", 257 } }); + run_dynamic_export_reuse_case( + device, "ggml_concat_dim0_f32", + { { "row_count", 3 }, { "lhs_span", 192 }, { "rhs_span", 96 } }, + { { "row_count", 7 }, { "lhs_span", 448 }, { "rhs_span", 224 } }, + { { "ggml.concat_dim0_f32.lhs_width", "64" }, { "ggml.concat_dim0_f32.rhs_width", "32" } }); + run_dynamic_export_reuse_case( + device, "ggml_rmsnorm_f32", { { "token_count", 1 } }, { { "token_count", 4 } }, + { { "ggml.rmsnorm_f32.hidden_size", "256" }, { "ggml.rmsnorm_f32.input_stride", "256" }, + { "ggml.rmsnorm_f32.rms_epsilon", "1e-06" } }); + run_dynamic_export_reuse_case( + device, "ggml_get_rows_f32", + { { "token_count", 3 }, { "row_count", 64 }, { "hidden_size", 256 } }, + { { "token_count", 7 }, { "row_count", 64 }, { "hidden_size", 256 } }, + { { "ggml.get_rows_f32.token_capacity", "8" }, { "ggml.get_rows_f32.hidden_capacity", "256" }, + { "ggml.get_rows_f32.weight_format", "32" } }); + run_dynamic_export_reuse_case( + device, "ggml_rope_f32", { { "token_count", 3 }, { "input_span", 768 } }, + { { "token_count", 7 }, { "input_span", 1792 } }, + { { "ggml.rope_f32.head_size", "64" }, { "ggml.rope_f32.n_dims", "64" }, + { "ggml.rope_f32.head_count", "4" }, { "ggml.rope_f32.token_capacity", "8" }, + { "ggml.rope_f32.input_stride1", "64" }, { "ggml.rope_f32.input_stride2", "256" }, + { "ggml.rope_f32.mscale", "1" }, { "ggml.rope_f32.mode", "0" } }); + run_dynamic_export_reuse_case( + device, "ggml_mul_mat_tiled_input_f32_publish_f32_aligned", { { "token_count", 32 } }, + { { "token_count", 64 } }, + { { "ggml.mul_mat.input_size", "256" }, { "ggml.mul_mat.output_size", "64" }, + { "ggml.mul_mat.output_accumulation", "0" }, { "ggml.mul_mat.output_unary_op", "23" }, + { "ggml.mul_mat.weight_format", "4" } }); + run_dynamic_export_reuse_case( + device, "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", { { "token_count", 32 } }, + { { "token_count", 64 } }, + { { "ggml.mul_mat_swiglu.input_size", "256" }, { "ggml.mul_mat_swiglu.output_size", "64" }, + { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, + { "ggml.mul_mat_swiglu.up_weight_format", "4" }, { "ggml.mul_mat_swiglu.op", "4" } }); + run_dynamic_export_reuse_case( + device, "llm_ssm_conv_f32", { { "n_t", 8 }, { "n_s", 1 } }, { { "n_t", 16 }, { "n_s", 2 } }, + { { "llm.ssm_conv.generic.d_conv", "4" }, { "llm.ssm_conv.generic.d_inner", "64" }, + { "llm.ssm_conv.generic.unary_op", "23" }, { "llm.ssm_conv.generic.workgroup_size", "256" } }); + run_dynamic_export_reuse_case( + device, "llm_ssm_conv_binary_f32", { { "n_t", 8 }, { "n_s", 1 } }, + { { "n_t", 16 }, { "n_s", 2 } }, + { { "llm.ssm_conv.generic.d_conv", "4" }, { "llm.ssm_conv.generic.d_inner", "64" }, + { "llm.ssm_conv.generic.unary_op", "23" }, { "llm.ssm_conv.generic.binary_op", "3" }, + { "llm.ssm_conv.generic.binary_lhs", "0" }, + { "llm.ssm_conv.generic.workgroup_size", "256" } }); + run_dynamic_export_reuse_case(device, "ggml_copy_transpose_f16", + { { "row_count", 512 }, { "column_count", 256 } }, + { { "row_count", 1024 }, { "column_count", 512 } }); + run_dynamic_export_reuse_case( + device, "ggml_flash_attention_f32_f16_wmma", flash_attention_workload(16, 64), + flash_attention_workload(17, 65), flash_attention_config("32", "4", "128", "0.0883883461")); + run_dynamic_export_reuse_case( + device, "ggml_flash_attention_f32_f16_wmma_publish_f16", flash_attention_workload(16, 64), + flash_attention_workload(17, 65), flash_attention_config("32", "4", "128", "0.0883883461")); + run_dynamic_export_reuse_case( + device, "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + flash_attention_decode_workload(65), flash_attention_decode_workload(127), + flash_attention_decode_config("32", "4", "128", "0.0883883461", "128")); + run_dynamic_export_reuse_case( + device, "llm_gated_delta_net_projection_epilogue_f32", { { "element_count", 48 } }, + { { "element_count", 192 } }, + { { "llm.gated_delta_net.epilogue_head_count", "48" }, + { "llm.gated_delta_net.epilogue_workgroup_size", "256" } }); run_targeted_export_materialize_case(device); } +static void run_binary_publication_compile_checks() { + std::string error; + std::unique_ptr jit = + ggml::hrx::create_loom_jit("gfx1100", ggml::hrx::LoomJitMode::Sync, error); + REQUIRE(jit != nullptr); + + const auto & binary_q8 = find_kernel("ggml_binary_f32_publish_q8_1_x4"); + require_compiled_kernel(compile_kernel(*jit, "binary-mul-q8", binary_q8, { + { "token_count", 2 }, + }, { + { "ggml.binary_f32.op", "2" }, + { "ggml.binary_f32.ne0", "2048" }, + { "ggml.binary_f32.src0_stride1", "2048" }, + { "ggml.binary_f32.src1_stride1", "2048" }, + { "ggml.binary_f32.src0_span", "4096" }, + { "ggml.binary_f32.src1_span", "4096" }, + })); + + const auto & binary_bc_q8 = find_kernel("ggml_binary_bc_f32_publish_q8_1_x4"); + require_compiled_kernel(compile_kernel(*jit, "binary-bc-mul-q8", binary_bc_q8, { + { "token_count", 2 }, + { "hidden_size", 256 }, + { "src0_element_count", 512 }, + { "src1_element_count", 256 }, + }, { + { "ggml.binary_bc_f32.op", "2" }, + { "ggml.binary_bc_f32.src0_broadcast_dim0", "0" }, + { "ggml.binary_bc_f32.src0_broadcast_dim1", "0" }, + { "ggml.binary_bc_f32.src1_broadcast_dim0", "0" }, + { "ggml.binary_bc_f32.src1_broadcast_dim1", "1" }, + })); +} + static bool has_hrx_test_device() { HrxTestDevice device; return device.open(); @@ -1334,6 +1545,7 @@ static bool has_hrx_test_device() { static void register_hrx_loom_jit_cases(test_runner::Suite & suite) { suite.host_case("compile", [] { run_loom_jit_compile_checks(); }); + suite.host_case("binary-publication", [] { run_binary_publication_compile_checks(); }); suite.host_case("moe-routing", [] { run_moe_routing_compile_checks(); }); suite.host_case("prefill-fusions", [] { run_prefill_fusion_compile_checks(); }); suite.host_case("binary-fusions", [] { run_binary_fusion_compile_checks(); }); diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 0d81b9c53eb2..12064f3f63c2 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -1,8 +1,10 @@ #include "backend-context.h" #include "dispatch/dispatch-scheduler.h" +#include "dispatch_registration/common/dispatch-activation-publication.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml-hrx.h" +#include "ggml-quants.h" #include "ggml.h" #include "graph/graph.h" #include "kernel-corpus/kernel-corpus.h" @@ -16,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -388,6 +391,22 @@ static std::vector scheduled_kernel_sequence(ggml_cgraph * graph) { return names; } +static ggml::hrx::KernelSpecialization scheduled_kernel_specialization(ggml_cgraph * graph, + const char * kernel_name) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + REQUIRE(scheduler.plan().valid()); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == kernel_name) { + return dispatch.kernel; + } + } + REQUIRE(false); + return {}; +} + static size_t producer_index_for_tensor(const ggml::hrx::Graph & graph, const ggml_tensor * tensor) { const ggml::hrx::Value * value = graph.values().find_tensor(tensor); REQUIRE(value != nullptr); @@ -438,7 +457,7 @@ static void match_dispatch_at_index(const ggml::hrx::Graph & graph, ggml::hrx::DispatchMatch & match) { REQUIRE(root_index < graph.nodes().size()); const ggml::hrx::DispatchMatchContext context = { - graph, &graph.nodes()[root_index], root_index, covered_nodes, plan, next_plan_value(graph, plan), + graph, &graph.nodes()[root_index], root_index, covered_nodes, plan, next_plan_value(graph, plan), ®istry, }; ggml::hrx::DispatchMatchDiagnostics diagnostics; if (!registry.match(context, match, &diagnostics)) { @@ -499,10 +518,11 @@ static void require_kernel_subsequence(const std::vector & sequence } static void run_alternate_value_alias_lookup_checks() { - constexpr int64_t element_count = 2048; - const size_t full_bytes = ggml_row_size(GGML_TYPE_F32, element_count); - const size_t f16_bytes = ggml_row_size(GGML_TYPE_F16, element_count); - const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, element_count); + constexpr int64_t element_count = 2048; + const size_t full_bytes = ggml_row_size(GGML_TYPE_F32, element_count); + const size_t f16_bytes = ggml_row_size(GGML_TYPE_F16, element_count); + const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, element_count); + const size_t symmetric_bytes = 1536; ggml::hrx::Graph graph; ggml::hrx::Status status; @@ -511,6 +531,10 @@ static void run_alternate_value_alias_lookup_checks() { const ggml::hrx::ValueId root(0); const ggml::hrx::ValueId full_alias(1); const ggml::hrx::ValueId partial_alias(2); + const ggml::hrx::ValueId reordered_alias(3); + const ggml::hrx::ValueId unrelated_same_storage(4); + const ggml::hrx::ValueId different_shape_alias(5); + const ggml::hrx::ValueId incompatible_q8_shape_alias(6); const ggml::hrx::ValueId f16_alternate(99); const ggml::hrx::ValueId q8_alternate(100); const ggml::hrx::ValueStorageId storage(0); @@ -526,9 +550,41 @@ static void run_alternate_value_alias_lookup_checks() { status = graph.values().add_snapshot_value( make_test_value(partial_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count / 2)); REQUIRE(status.success()); + ggml::hrx::Value reordered = + make_test_value(reordered_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); + reordered.contiguous = false; + status = graph.values().add_snapshot_value(std::move(reordered)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_test_value(unrelated_same_storage, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count)); + REQUIRE(status.success()); + ggml::hrx::Value different_shape = + make_test_value(different_shape_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); + different_shape.ne = { element_count / 2, 2, 1, 1 }; + different_shape.nb = { sizeof(float), static_cast(element_count / 2) * sizeof(float), full_bytes, + full_bytes }; + status = graph.values().add_snapshot_value(std::move(different_shape)); + REQUIRE(status.success()); + ggml::hrx::Value incompatible_q8_shape = + make_test_value(incompatible_q8_shape_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); + incompatible_q8_shape.ne = { 64, element_count / 64, 1, 1 }; + incompatible_q8_shape.nb = { sizeof(float), 64 * sizeof(float), full_bytes, full_bytes }; + status = graph.values().add_snapshot_value(std::move(incompatible_q8_shape)); + REQUIRE(status.success()); + + graph.add_node(GGML_OP_RESHAPE, full_alias, { root }); + graph.add_node(GGML_OP_VIEW, partial_alias, { root }); + graph.add_node(GGML_OP_TRANSPOSE, reordered_alias, { root }); + graph.add_node(GGML_OP_CPY, unrelated_same_storage, { root }); + graph.add_node(GGML_OP_RESHAPE, different_shape_alias, { root }); + graph.add_node(GGML_OP_RESHAPE, incompatible_q8_shape_alias, { root }); + status = graph.build_index(); + REQUIRE(status.success()); REQUIRE(plan.metadata.append_alternate_value({ root, q8_alternate, GGML_TYPE_Q8_1, q8_bytes, "q8" }, status)); REQUIRE(plan.metadata.append_alternate_value({ root, f16_alternate, GGML_TYPE_F16, f16_bytes, "f16" }, status)); + REQUIRE(plan.metadata.append_alternate_value( + { root, ggml::hrx::ValueId(104), GGML_TYPE_COUNT, symmetric_bytes, "symmetric-i4" }, status)); const ggml::hrx::CommandPlanAlternateValue * exact = ggml::hrx::find_alternate_value(graph, plan, root, GGML_TYPE_Q8_1, q8_bytes); @@ -544,14 +600,412 @@ static void run_alternate_value_alias_lookup_checks() { ggml::hrx::find_alternate_value(graph, plan, partial_alias, GGML_TYPE_Q8_1, q8_bytes); REQUIRE(through_partial_alias == nullptr); + const ggml::hrx::CommandPlanAlternateValue * through_reordered_alias = + ggml::hrx::find_alternate_value(graph, plan, reordered_alias, GGML_TYPE_Q8_1, q8_bytes); + REQUIRE(through_reordered_alias == nullptr); + + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, unrelated_same_storage, GGML_TYPE_Q8_1, q8_bytes) == nullptr); + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_Q8_1, q8_bytes) != nullptr); + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, incompatible_q8_shape_alias, GGML_TYPE_Q8_1, q8_bytes) == + nullptr); + const ggml::hrx::CommandPlanAlternateValue * f16_exact = ggml::hrx::find_alternate_value(graph, plan, root, GGML_TYPE_F16, f16_bytes); REQUIRE(f16_exact != nullptr); REQUIRE(f16_exact->alternate_value == f16_alternate); + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_F16, f16_bytes) != nullptr); + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, full_alias, GGML_TYPE_COUNT, symmetric_bytes) != nullptr); + REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_COUNT, symmetric_bytes) == + nullptr); ggml::hrx::Status conflict_status; REQUIRE(!plan.metadata.append_alternate_value( { root, ggml::hrx::ValueId(101), GGML_TYPE_Q8_1, q8_bytes, "q8-conflict" }, conflict_status)); + + ggml::hrx::Status generated_status; + const ggml::hrx::ValueId k16_value(102); + REQUIRE(plan.metadata.append_generated_resource( + { root, ggml::hrx::GeneratedResourceRole::F16K16Major, k16_value, f16_bytes, {} }, generated_status)); + const ggml::hrx::CommandPlanGeneratedResource * generated_through_alias = ggml::hrx::find_generated_resource( + graph, plan, full_alias, ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes); + REQUIRE(generated_through_alias != nullptr); + REQUIRE(generated_through_alias->generated_value == k16_value); + REQUIRE(ggml::hrx::find_generated_resource(graph, plan, partial_alias, + ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); + REQUIRE(ggml::hrx::find_generated_resource(graph, plan, reordered_alias, + ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); + REQUIRE(ggml::hrx::find_generated_resource(graph, plan, different_shape_alias, + ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); +} + +static bool publication_accepts_add(const ggml::hrx::DispatchMatchContext &, + const ggml::hrx::GraphNode & consumer, + const ggml::hrx::Value & input) { + return consumer.op == GGML_OP_ADD && + std::find(consumer.inputs.begin(), consumer.inputs.end(), input.id) != consumer.inputs.end(); +} + +static bool publication_accepts_mul(const ggml::hrx::DispatchMatchContext &, + const ggml::hrx::GraphNode & consumer, + const ggml::hrx::Value & input) { + return consumer.op == GGML_OP_MUL && + std::find(consumer.inputs.begin(), consumer.inputs.end(), input.id) != consumer.inputs.end(); +} + +static bool publication_accepts_mul_output_13(const ggml::hrx::DispatchMatchContext & context, + const ggml::hrx::GraphNode & consumer, + const ggml::hrx::Value & input) { + return consumer.output == ggml::hrx::ValueId(13) && publication_accepts_mul(context, consumer, input); +} + +static bool publication_accepts_mul_output_15(const ggml::hrx::DispatchMatchContext & context, + const ggml::hrx::GraphNode & consumer, + const ggml::hrx::Value & input) { + return consumer.output == ggml::hrx::ValueId(15) && publication_accepts_mul(context, consumer, input); +} + +static void run_activation_publication_contract_checks() { + constexpr int64_t row_size = 2048; + constexpr int64_t row_count = 2; + constexpr int64_t elements = row_size * row_count; + const size_t f32_bytes = ggml_row_size(GGML_TYPE_F32, elements); + + auto make_activation = [&](ggml::hrx::ValueId id, ggml::hrx::ValueStorageId storage, ggml::hrx::ValueId root, + ggml::hrx::ValueId alias, int64_t element_count = elements) { + ggml::hrx::Value value = make_test_value(id, storage, root, alias, 0, f32_bytes, GGML_TYPE_F32, element_count); + if (element_count == elements) { + value.ne = { row_size, row_count, 1, 1 }; + value.nb = { sizeof(float), static_cast(row_size) * sizeof(float), f32_bytes, f32_bytes }; + } + return value; + }; + + ggml::hrx::Graph graph; + ggml::hrx::Status status; + const ggml::hrx::ValueId root(0); + const ggml::hrx::ValueId reshaped(1); + const ggml::hrx::ValueId partial(2); + const ggml::hrx::ValueId transposed(3); + const ggml::hrx::ValueId different_shape(8); + const ggml::hrx::ValueId incompatible_q8_shape(10); + const ggml::hrx::ValueId shape_back(12); + const ggml::hrx::ValueId mismatched_alias_source(14); + status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(0), root, f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_activation(root, ggml::hrx::ValueStorageId(0), root, ggml::hrx::ValueId())); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(reshaped, ggml::hrx::ValueStorageId(0), root, root)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_activation(partial, ggml::hrx::ValueStorageId(0), root, root, elements / 2)); + REQUIRE(status.success()); + ggml::hrx::Value reordered = make_activation(transposed, ggml::hrx::ValueStorageId(0), root, root); + reordered.contiguous = false; + status = graph.values().add_snapshot_value(std::move(reordered)); + REQUIRE(status.success()); + + for (int32_t value = 4; value < 8; ++value) { + const ggml::hrx::ValueStorageId storage(value - 3); + const ggml::hrx::ValueId id(value); + status = graph.values().add_snapshot_storage({ storage, id, f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(id, storage, id, ggml::hrx::ValueId())); + REQUIRE(status.success()); + } + ggml::hrx::Value different = make_activation(different_shape, ggml::hrx::ValueStorageId(0), root, root); + different.ne = { row_size / 2, row_count * 2, 1, 1 }; + different.nb = { sizeof(float), static_cast(row_size / 2) * sizeof(float), f32_bytes, f32_bytes }; + status = graph.values().add_snapshot_value(std::move(different)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(5), ggml::hrx::ValueId(9), f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(9), ggml::hrx::ValueStorageId(5), + ggml::hrx::ValueId(9), ggml::hrx::ValueId())); + REQUIRE(status.success()); + ggml::hrx::Value incompatible = make_activation(incompatible_q8_shape, ggml::hrx::ValueStorageId(0), root, root); + incompatible.ne = { 64, elements / 64, 1, 1 }; + incompatible.nb = { sizeof(float), 64 * sizeof(float), f32_bytes, f32_bytes }; + status = graph.values().add_snapshot_value(std::move(incompatible)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(6), ggml::hrx::ValueId(11), f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(11), ggml::hrx::ValueStorageId(6), + ggml::hrx::ValueId(11), ggml::hrx::ValueId())); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value( + make_activation(shape_back, ggml::hrx::ValueStorageId(0), root, different_shape)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(7), ggml::hrx::ValueId(13), f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(13), ggml::hrx::ValueStorageId(7), + ggml::hrx::ValueId(13), ggml::hrx::ValueId())); + REQUIRE(status.success()); + ggml::hrx::Value mismatched = make_activation( + mismatched_alias_source, ggml::hrx::ValueStorageId(0), root, root); + mismatched.ne = { row_size / 2, row_count * 2, 1, 1 }; + mismatched.nb = { sizeof(float), static_cast(row_size / 2) * sizeof(float), f32_bytes, f32_bytes }; + status = graph.values().add_snapshot_value(std::move(mismatched)); + REQUIRE(status.success()); + status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(8), ggml::hrx::ValueId(15), f32_bytes }); + REQUIRE(status.success()); + status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(15), ggml::hrx::ValueStorageId(8), + ggml::hrx::ValueId(15), ggml::hrx::ValueId())); + REQUIRE(status.success()); + + graph.add_node(GGML_OP_RESHAPE, reshaped, { root }); + graph.add_node(GGML_OP_VIEW, partial, { root }); + graph.add_node(GGML_OP_TRANSPOSE, transposed, { root }); + graph.add_node(GGML_OP_ADD, ggml::hrx::ValueId(4), { root, root }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(5), { reshaped, reshaped }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(6), { partial, partial }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(7), { transposed, transposed }); + graph.add_node(GGML_OP_RESHAPE, different_shape, { root }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(9), { different_shape, different_shape }); + graph.add_node(GGML_OP_RESHAPE, incompatible_q8_shape, { root }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(11), { incompatible_q8_shape, incompatible_q8_shape }); + graph.add_node(GGML_OP_RESHAPE, shape_back, { different_shape }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(13), { shape_back, shape_back }); + graph.add_node(GGML_OP_RESHAPE, mismatched_alias_source, { different_shape }); + graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(15), { mismatched_alias_source, mismatched_alias_source }); + status = graph.build_index(); + REQUIRE(status.success()); + + ggml::hrx::DispatchRegistryBuilder publication_registry_builder; + publication_registry_builder.add_activation_consumer({ + "test.q8.mul", + GGML_OP_MUL, + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_mul, + }); + publication_registry_builder.add_activation_consumer({ + "test.f16-row.add", + GGML_OP_ADD, + ggml::hrx::DispatchActivationInputF16Row, + ggml::hrx::DispatchActivationConsumerUse::Row, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_add, + }); + const ggml::hrx::DispatchRegistry publication_registry = publication_registry_builder.build(); + const std::vector covered(graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + const ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &publication_registry, + }; + const ggml::hrx::Value * produced = graph.values().find(root); + REQUIRE(produced != nullptr); + + const ggml::hrx::CommonActivationPublicationCapabilities direct_only = { + ggml::hrx::DispatchActivationInputQ8_1X4, + }; + const ggml::hrx::CommonActivationPublicationPlan direct_only_plan = + ggml::hrx::common_select_activation_publication_plan(context, *produced, direct_only); + REQUIRE(!direct_only_plan.matched()); + REQUIRE(direct_only_plan.fallback_reason == + ggml::hrx::CommonActivationPublicationFallbackReason::NoQualifiedConsumer); + + const ggml::hrx::CommonActivationPublicationCapabilities q8_capabilities = { + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationInputQ8_1X4, + }; + const ggml::hrx::CommonActivationPublicationPlan q8_plan = + ggml::hrx::common_select_activation_publication_plan(context, *produced, q8_capabilities); + REQUIRE(q8_plan.matched()); + const ggml::hrx::CommonActivationPublicationDemand q8_demand = q8_plan.preferred(); + REQUIRE(q8_demand.format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); + REQUIRE(q8_demand.consumer_value->id == reshaped); + REQUIRE(q8_demand.alias_depth == 1); + REQUIRE(q8_demand.mixed_fanout); + + const ggml::hrx::CommonActivationPublicationCapabilities f16_capabilities = { + ggml::hrx::DispatchActivationInputF16Row, + }; + const ggml::hrx::CommonActivationPublicationPlan f16_plan = + ggml::hrx::common_select_activation_publication_plan(context, *produced, f16_capabilities); + REQUIRE(f16_plan.matched()); + const ggml::hrx::CommonActivationPublicationDemand f16_demand = f16_plan.preferred(); + REQUIRE(f16_demand.format == ggml::hrx::CommonActivationPublicationFormat::F16Row); + REQUIRE(f16_demand.consumer_value->id == root); + REQUIRE(f16_demand.alias_depth == 0); + REQUIRE(f16_demand.mixed_fanout); + + ggml::hrx::DispatchRegistryBuilder direct_k16_registry_builder; + direct_k16_registry_builder.add_activation_consumer({ + "test.k16.mul", + GGML_OP_MUL, + ggml::hrx::DispatchActivationInputF16K16Major, + ggml::hrx::DispatchActivationConsumerUse::Tiled, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_mul, + }); + const ggml::hrx::DispatchRegistry direct_k16_registry = direct_k16_registry_builder.build(); + const ggml::hrx::DispatchMatchContext direct_k16_context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &direct_k16_registry, + }; + const ggml::hrx::CommonActivationPublicationCapabilities mixed_alias_capabilities = { + ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16K16Major, + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationProducerRequirementNone, + }; + const ggml::hrx::CommonActivationPublicationDemand direct_k16_through_q8_alias = + ggml::hrx::common_select_activation_publication_plan( + direct_k16_context, *produced, mixed_alias_capabilities).publications.front(); + REQUIRE(!direct_k16_through_q8_alias.matched()); + + ggml::hrx::DispatchRegistryBuilder shape_back_registry_builder; + shape_back_registry_builder.add_activation_consumer({ + "test.k16.shape-back", + GGML_OP_MUL, + ggml::hrx::DispatchActivationInputF16K16Major, + ggml::hrx::DispatchActivationConsumerUse::Tiled, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_mul_output_13, + }); + const ggml::hrx::DispatchRegistry shape_back_registry = shape_back_registry_builder.build(); + const ggml::hrx::DispatchMatchContext shape_back_context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &shape_back_registry, + }; + REQUIRE(!ggml::hrx::common_select_activation_publication_plan( + shape_back_context, *produced, mixed_alias_capabilities).matched()); + + ggml::hrx::CommonActivationPublicationCapabilities multiple_capabilities = { + ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16Row, + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16Row, + }; + ggml::hrx::DispatchRegistryBuilder mismatched_alias_registry_builder; + mismatched_alias_registry_builder.add_activation_consumer({ + "test.q8.mismatched-alias-source", + GGML_OP_MUL, + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_mul_output_15, + }); + const ggml::hrx::DispatchRegistry mismatched_alias_registry = mismatched_alias_registry_builder.build(); + const ggml::hrx::DispatchMatchContext mismatched_alias_context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &mismatched_alias_registry, + }; + REQUIRE(!ggml::hrx::common_select_activation_publication_plan( + mismatched_alias_context, *produced, multiple_capabilities).matched()); + + const ggml::hrx::CommonActivationPublicationPlan multiple_plan = + ggml::hrx::common_select_activation_publication_plan(context, *produced, multiple_capabilities); + REQUIRE(multiple_plan.publication_count == 2); + REQUIRE(multiple_plan.publications[0].format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); + REQUIRE(multiple_plan.publications[1].format == ggml::hrx::CommonActivationPublicationFormat::F16Row); + ggml::hrx::DispatchRegistryBuilder q8_only_registry_builder; + q8_only_registry_builder.add_activation_consumer({ + "test.q8.only", + GGML_OP_MUL, + ggml::hrx::DispatchActivationInputQ8_1X4, + ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, + ggml::hrx::DispatchActivationProducerRequirementNone, + ggml::hrx::DispatchSource::Common, + publication_accepts_mul, + }); + const ggml::hrx::DispatchRegistry q8_only_registry = q8_only_registry_builder.build(); + const ggml::hrx::DispatchMatchContext q8_only_context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &q8_only_registry, + }; + const ggml::hrx::CommonActivationPublicationPlan companion_plan = + ggml::hrx::common_select_activation_publication_plan(q8_only_context, *produced, multiple_capabilities); + REQUIRE(companion_plan.publication_count == 2); + REQUIRE(companion_plan.publications[0].format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); + REQUIRE(companion_plan.publications[1].format == ggml::hrx::CommonActivationPublicationFormat::F16Row); + REQUIRE(companion_plan.publications[1].consumer == companion_plan.publications[0].consumer); + multiple_capabilities.co_publication_formats = ggml::hrx::DispatchActivationInputNone; + REQUIRE(ggml::hrx::common_select_activation_publication_plan(context, *produced, multiple_capabilities) + .publication_count == 1); + + const ggml::hrx::CommonActivationPublicationCapabilities disabled_capabilities; + const ggml::hrx::CommonActivationPublicationPlan disabled_plan = + ggml::hrx::common_select_activation_publication_plan(context, *produced, disabled_capabilities); + REQUIRE(!disabled_plan.matched()); + REQUIRE(disabled_plan.fallback_reason == + ggml::hrx::CommonActivationPublicationFallbackReason::NoEnabledCandidates); + + ggml::hrx::DispatchMatch match; + match.completion_counter_requests.push_back({ ggml::hrx::ValueId(100), "existing.counter", 1 }); + const ggml::hrx::CommonActivationPublicationDemand symmetric_demand = { + ggml::hrx::CommonActivationPublicationFormat::SymmetricI4K32, + produced, + &graph.nodes()[3], + }; + ggml::hrx::CommonActivationPublication symmetric; + ggml::hrx::CommonActivationPublication q8; + ggml::hrx::CommonActivationPublication f16; + ggml::hrx::CommonActivationPublication k16; + REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, symmetric_demand, + "test.symmetric_i4", symmetric)); + REQUIRE(ggml::hrx::common_append_activation_publication(match, symmetric, "test.symmetric_i4")); + REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, q8_demand, "test.q8", q8)); + REQUIRE(ggml::hrx::common_append_activation_publication(match, q8, "test.q8")); + REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, f16_demand, "test.f16", f16)); + REQUIRE(ggml::hrx::common_append_activation_publication(match, f16, "test.f16")); + const ggml::hrx::CommonActivationPublicationDemand k16_demand = { + ggml::hrx::CommonActivationPublicationFormat::F16K16Major, + q8_demand.consumer_value, + q8_demand.consumer, + }; + REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, k16_demand, "test.k16", k16)); + REQUIRE(ggml::hrx::common_append_activation_publication(match, k16, "test.k16")); + + ggml::hrx::DispatchMatch invalid_id_match; + const ggml::hrx::DispatchMatchContext invalid_id_context = { + graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(-1), + }; + ggml::hrx::CommonActivationPublication invalid_id_publication; + REQUIRE(!ggml::hrx::common_reserve_activation_publication(invalid_id_context, invalid_id_match, *produced, + f16_demand, "test.invalid_id", invalid_id_publication)); + REQUIRE(invalid_id_match.transients.empty()); + + ggml::hrx::DispatchMatch conflict_match; + REQUIRE(ggml::hrx::common_append_activation_publication(conflict_match, q8, "test.q8")); + ggml::hrx::CommonActivationPublication conflicting_q8 = q8; + conflicting_q8.alternate_value = ggml::hrx::ValueId(200); + REQUIRE(!ggml::hrx::common_append_activation_publication(conflict_match, conflicting_q8, "test.q8.conflict")); + REQUIRE(!conflict_match.status.success()); + + REQUIRE(match.status.success()); + REQUIRE(match.transients.size() == 4); + REQUIRE(symmetric.alternate_value == ggml::hrx::ValueId(101)); + REQUIRE(q8.alternate_value == ggml::hrx::ValueId(102)); + REQUIRE(f16.alternate_value == ggml::hrx::ValueId(103)); + REQUIRE(k16.alternate_value == ggml::hrx::ValueId(104)); + REQUIRE(symmetric.payload_bytes == static_cast(elements) / 2); + REQUIRE(symmetric.scales_offset % 256 == 0); + REQUIRE(symmetric.sums_offset % 256 == 0); + REQUIRE(q8.byte_count == static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, row_size)); + REQUIRE(f16.byte_count == static_cast(elements) * sizeof(ggml_fp16_t)); + REQUIRE(symmetric.binding(symmetric.scales_offset, symmetric.scales_bytes).value == symmetric.alternate_value); + + ggml::hrx::CommandPlan recorded; + recorded.metadata = std::move(match.metadata); + const std::vector & diagnostics = + recorded.metadata.activation_publication_diagnostics(); + REQUIRE(diagnostics.size() == 4); + REQUIRE(diagnostics[1].source_value == root); + REQUIRE(diagnostics[1].alternate_value == q8.alternate_value); + REQUIRE(diagnostics[1].consumer_value == reshaped); + REQUIRE(diagnostics[1].requested_format == "q8-1-x4"); + REQUIRE(diagnostics[1].publication_name == "test.q8"); + REQUIRE(diagnostics[1].publication_stage == "published"); + REQUIRE(diagnostics[1].fallback_reason == "none"); + REQUIRE(diagnostics[1].alias_depth == 1); + REQUIRE(diagnostics[1].mixed_fanout); + REQUIRE(ggml::hrx::find_alternate_value(graph, recorded, reshaped, GGML_TYPE_Q8_1, q8.byte_count) != nullptr); + REQUIRE(ggml::hrx::find_alternate_value(graph, recorded, reshaped, GGML_TYPE_F16, f16.byte_count) != nullptr); + REQUIRE(ggml::hrx::find_generated_resource(graph, recorded, reshaped, ggml::hrx::GeneratedResourceRole::F16K16Major, + k16.byte_count) != nullptr); } static std::vector make_pattern_f32(size_t element_count, int seed, float scale = 0.01f) { @@ -600,10 +1054,27 @@ static std::vector make_pattern_f16(size_t element_count, int seed, static std::vector make_quantized_rows(ggml_type type, int64_t row_length, int64_t row_count, int seed) { const ggml_type_traits * traits = ggml_get_type_traits(type); REQUIRE(traits != nullptr); - REQUIRE(traits->from_float_ref != nullptr); ggml_quantize_init(type); const size_t row_size = ggml_row_size(type, row_length); std::vector data(static_cast(row_count) * row_size); + if (type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS) { + std::vector source(static_cast(row_length * row_count)); + std::vector importance(static_cast(row_length), 1.0f); + for (int64_t r = 0; r < row_count; ++r) { + for (int64_t c = 0; c < row_length; ++c) { + const int value = static_cast((r * 13 + c * 7 + seed * 31) % 101) - 50; + source[static_cast(r * row_length + c)] = static_cast(value) * 0.005f; + } + } + const size_t bytes_written = type == GGML_TYPE_IQ2_XXS ? + quantize_iq2_xxs(source.data(), data.data(), row_count, row_length, + importance.data()) : + quantize_iq2_xs(source.data(), data.data(), row_count, row_length, + importance.data()); + REQUIRE(bytes_written == data.size()); + return data; + } + REQUIRE(traits->from_float_ref != nullptr); std::vector row(static_cast(row_length)); for (int64_t r = 0; r < row_count; ++r) { for (int64_t c = 0; c < row_length; ++c) { @@ -615,7 +1086,59 @@ static std::vector make_quantized_rows(ggml_type type, int64_t row_leng return data; } +static std::vector make_iq1_rows(ggml_type type, int64_t row_length, int64_t row_count, int seed) { + REQUIRE(type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M); + REQUIRE(row_length % 256 == 0); + const size_t row_size = ggml_row_size(type, row_length); + const size_t block_size = type == GGML_TYPE_IQ1_S ? 50 : 56; + REQUIRE(row_size == static_cast(row_length / 256) * block_size); + std::vector data(static_cast(row_count) * row_size); + for (int64_t row = 0; row < row_count; ++row) { + uint8_t * row_data = data.data() + static_cast(row) * row_size; + for (int64_t block = 0; block < row_length / 256; ++block) { + uint8_t * dst = row_data + static_cast(block) * block_size; + if (type == GGML_TYPE_IQ1_S) { + const ggml_fp16_t d = ggml_fp32_to_fp16(0.015625f); + std::memcpy(dst, &d, sizeof(d)); + for (int i = 0; i < 32; ++i) { + dst[2 + i] = static_cast((row * 13 + block * 29 + i * 37 + seed) & 0xff); + } + for (int group = 0; group < 8; ++group) { + uint16_t qh = 0; + for (int entry = 0; entry < 4; ++entry) { + qh |= static_cast(((row + block + group + entry + seed) & 7) << (3 * entry)); + } + qh |= static_cast(((row + block + group) & 7) << 12); + qh |= static_cast(((row + group) & 1) << 15); + std::memcpy(dst + 34 + 2 * group, &qh, sizeof(qh)); + } + } else { + for (int i = 0; i < 32; ++i) { + dst[i] = static_cast((row * 17 + block * 31 + i * 41 + seed) & 0xff); + } + for (int i = 0; i < 16; ++i) { + const uint8_t low = static_cast((row + block + i + seed) & 7); + const uint8_t high = static_cast((row + block + i + seed + 3) & 7); + dst[32 + i] = static_cast(low | ((i & 1) ? 0x08 : 0) | (high << 4) | + ((i & 2) ? 0x80 : 0)); + } + const uint16_t base_bits = ggml_fp32_to_fp16(0.015625f); + for (int word = 0; word < 4; ++word) { + const uint16_t local_scales = static_cast(1 | (2 << 3) | (3 << 6) | (4 << 9)); + const uint16_t base_nibble = static_cast((base_bits >> (4 * word)) & 0xf); + const uint16_t scales = static_cast(local_scales | (base_nibble << 12)); + std::memcpy(dst + 48 + 2 * word, &scales, sizeof(scales)); + } + } + } + } + return data; +} + static std::vector make_matmul_weight_bytes(ggml_type type, int64_t row_length, int64_t row_count, int seed) { + if (type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M) { + return make_iq1_rows(type, row_length, row_count, seed); + } if (type == GGML_TYPE_F32) { (void) seed; const std::vector weights(static_cast(row_length * row_count), 0.00390625f); @@ -1385,7 +1908,7 @@ static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, }; Graph a = build(cpu_ctx); Graph b = build(hrx_ctx); - require_kernel_subsequence(scheduled_kernel_sequence(b.graph), { "loom_libs:ggml_rmsnorm_gate_f32_f16" }); + require_kernel_subsequence(scheduled_kernel_sequence(b.graph), { "loom_libs:ggml_rmsnorm_gate_f32_publish" }); ggml_backend_buffer_t a_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); ggml_backend_buffer_t b_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx); REQUIRE(a_buffer != nullptr && b_buffer != nullptr); @@ -1457,6 +1980,8 @@ static void run_rmsnorm_cpu_reference_case(int64_t hidden_size = 256, ggml_build_forward_expand(hrx_graph, hrx_output); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_f32" }); + REQUIRE(scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_rmsnorm_f32").workload_specialization == + ggml::hrx::WorkloadSpecialization::Dynamic); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -1892,6 +2417,10 @@ static void run_scheduled_hrx_scale_rmsnorm_case() { REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); + const ggml::hrx::KernelSpecialization scale_specialization = + scheduled_kernel_specialization(graph, "loom_libs:ggml_scale_f32"); + REQUIRE(scale_specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); REQUIRE(ggml_backend_sched_get_tensor_backend(sched, scaled) == hrx_backend); REQUIRE(ggml_backend_sched_get_tensor_backend(sched, output) == hrx_backend); @@ -2259,6 +2788,15 @@ static void run_qwen_flash_attention_case() { require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8" }); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); + REQUIRE(output_value != nullptr); + REQUIRE(scheduler.plan().metadata.find_alternate_value( + output_value->id, GGML_TYPE_Q8_1, + static_cast(kQueryTokenCount) * ggml_row_size(GGML_TYPE_Q8_1, kQwenFlashHeadSize)) == nullptr); static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); const bool had_original_env = original_env != nullptr; @@ -2350,6 +2888,11 @@ 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); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -2420,6 +2963,8 @@ 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); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -2549,7 +3094,10 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ const int64_t output_hidden_size = query_head_count * value_head_size; const size_t q8_row_bytes = ggml_row_size(GGML_TYPE_Q8_1, output_hidden_size); const size_t q8_output_bytes = static_cast(query_token_count) * q8_row_bytes; - const int64_t key_value_capacity = (key_value_token_count + 63) / 64 * 64; + int64_t key_value_capacity = 64; + while (key_value_capacity < key_value_token_count) { + key_value_capacity *= 2; + } const int64_t key_value_blocks = key_value_capacity / 64; const size_t partial_scalar_bytes = static_cast(key_value_head_count * key_value_blocks * 16) * sizeof(float); @@ -2564,10 +3112,18 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ QwenFlashAttentionLayoutGraph attention = build_qwen_flash_attention_layout_graph( ctx, query_token_count, key_value_token_count, qk_head_size, value_head_size, scale); + ggml_tensor * attention_output = + ggml_reshape_2d(ctx, attention.output, output_hidden_size, query_token_count); + ggml_tensor * consumer_weight = + ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, output_hidden_size, kQwenHiddenSize); + ggml_tensor * consumer_output = ggml_mul_mat(ctx, consumer_weight, attention_output); + REQUIRE(attention_output != nullptr); + REQUIRE(consumer_weight != nullptr); + REQUIRE(consumer_output != nullptr); ggml_cgraph * cgraph = ggml_new_graph(ctx); REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, attention.output); + ggml_build_forward_expand(cgraph, consumer_output); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); REQUIRE(imported.valid()); @@ -2576,7 +3132,7 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); const ggml::hrx::CommandPlan & plan = scheduler.plan(); REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == static_cast(query_token_count)); + REQUIRE(plan.dispatches.size() == static_cast(query_token_count + 1)); REQUIRE(plan.transients.size() == 4); REQUIRE(plan.completion_counter_requests.size() == 1); @@ -2586,14 +3142,16 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ REQUIRE(plan.transients[3].size == q8_output_bytes); REQUIRE(plan.completion_counter_requests[0].count == key_value_head_count); - const ggml::hrx::Value * query_value = imported.graph.values().find_tensor(attention.query); - const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.mask); - const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(attention.output); + const ggml::hrx::Value * query_value = imported.graph.values().find_tensor(attention.query); + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.mask); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(attention.output); + const ggml::hrx::Value * publication_value = imported.graph.values().find_tensor(attention_output); REQUIRE(query_value != nullptr); REQUIRE(mask_value != nullptr); REQUIRE(output_value != nullptr); + REQUIRE(publication_value != nullptr); const ggml::hrx::CommandPlanAlternateValue * q8_alternate = - plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_output_bytes); + plan.metadata.find_alternate_value(publication_value->id, GGML_TYPE_Q8_1, q8_output_bytes); REQUIRE(q8_alternate != nullptr); REQUIRE(q8_alternate->alternate_value == plan.transients[3].value); @@ -2604,6 +3162,7 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ const ggml::hrx::Dispatch & dispatch = plan.dispatches[static_cast(row)]; const char * expected_kernel = "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"; REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel); + REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.query_head_count") == std::to_string(query_head_count)); @@ -2670,6 +3229,8 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); ggml_tensor * rms = ggml_rms_norm(ctx, residual, kQwenRmsNormEps); ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); + ggml_tensor * next_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenHiddenSize); + ggml_tensor * next_output = ggml_mul_mat(ctx, next_weight, normalized); REQUIRE(attention_output != nullptr); REQUIRE(output_weight != nullptr); REQUIRE(projection != nullptr); @@ -2678,10 +3239,12 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(norm_weight != nullptr); REQUIRE(rms != nullptr); REQUIRE(normalized != nullptr); + REQUIRE(next_weight != nullptr); + REQUIRE(next_output != nullptr); ggml_cgraph * cgraph = ggml_new_graph(ctx); REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, normalized); + ggml_build_forward_expand(cgraph, next_output); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); REQUIRE(imported.valid()); @@ -2690,10 +3253,15 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); const ggml::hrx::CommandPlan & plan = scheduler.plan(); REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 2); + REQUIRE(plan.dispatches.size() == 3); REQUIRE(plan.completion_counter_requests.size() == 2); - const ggml::hrx::Dispatch & projection_dispatch = plan.dispatches.back(); + const auto projection_it = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"; + }); + REQUIRE(projection_it != plan.dispatches.end()); + const ggml::hrx::Dispatch & projection_dispatch = *projection_it; REQUIRE(kernel_name_for_id(projection_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"); REQUIRE(projection_dispatch.kernel.integer_parameters.at("token_count") == query_token_count); @@ -2704,7 +3272,7 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_accumulation") == "1"); REQUIRE(projection_dispatch.bindings.size() == 7); - const ggml::hrx::Value * flash_output_value = imported.graph.values().find_tensor(attention.output); + const ggml::hrx::Value * flash_output_value = imported.graph.values().find_tensor(attention_output); const ggml::hrx::Value * residual_input_value = imported.graph.values().find_tensor(residual_input); const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); @@ -2713,9 +3281,10 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(residual_value != nullptr); REQUIRE(normalized_value != nullptr); const ggml::hrx::CommandPlanAlternateValue * attention_q8 = - ggml::hrx::find_alternate_value(plan, flash_output_value->id, GGML_TYPE_Q8_1, attention_q8_bytes); + ggml::hrx::find_alternate_value(imported.graph, plan, flash_output_value->id, GGML_TYPE_Q8_1, + attention_q8_bytes); const ggml::hrx::CommandPlanAlternateValue * next_q8 = - ggml::hrx::find_alternate_value(plan, normalized_value->id, GGML_TYPE_Q8_1, next_q8_bytes); + ggml::hrx::find_alternate_value(imported.graph, plan, normalized_value->id, GGML_TYPE_Q8_1, next_q8_bytes); REQUIRE(attention_q8 != nullptr); REQUIRE(next_q8 != nullptr); REQUIRE(projection_dispatch.bindings[0].value == attention_q8->alternate_value); @@ -2787,6 +3356,86 @@ static void run_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_dynamic_generic_ops_cpu_reference_case() { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t element_count = 257; + ggml_tensor * cpu_input = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_input = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_unary = ggml_sqr(cpu_ctx, cpu_input); + ggml_tensor * hrx_unary = ggml_sqr(hrx_ctx, hrx_input); + ggml_tensor * cpu_copy_target = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_copy_target = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_copy = ggml_cpy(cpu_ctx, cpu_input, cpu_copy_target); + ggml_tensor * hrx_copy = ggml_cpy(hrx_ctx, hrx_input, hrx_copy_target); + ggml_tensor * cpu_concat_lhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 64, 7); + ggml_tensor * cpu_concat_rhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 32, 7); + ggml_tensor * hrx_concat_lhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 64, 7); + ggml_tensor * hrx_concat_rhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 32, 7); + ggml_tensor * cpu_concat = ggml_concat(cpu_ctx, cpu_concat_lhs, cpu_concat_rhs, 0); + ggml_tensor * hrx_concat = ggml_concat(hrx_ctx, hrx_concat_lhs, hrx_concat_rhs, 0); + REQUIRE(cpu_unary != nullptr && hrx_unary != nullptr); + REQUIRE(cpu_copy != nullptr && hrx_copy != nullptr); + REQUIRE(cpu_concat != nullptr && hrx_concat != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr && hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_unary); + ggml_build_forward_expand(cpu_graph, cpu_copy); + ggml_build_forward_expand(cpu_graph, cpu_concat); + ggml_build_forward_expand(hrx_graph, hrx_unary); + ggml_build_forward_expand(hrx_graph, hrx_copy); + ggml_build_forward_expand(hrx_graph, hrx_concat); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", + "loom_libs:ggml_concat_dim0_f32" }); + for (const char * kernel_name : + { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", "loom_libs:ggml_concat_dim0_f32" }) { + const ggml::hrx::KernelSpecialization specialization = + scheduled_kernel_specialization(hrx_graph, kernel_name); + REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + } + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr && hrx_buffer != nullptr); + const std::vector input = make_pattern_f32(element_count, 73, 0.03125f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + const std::vector concat_lhs = make_pattern_f32(64 * 7, 79, 0.015625f); + const std::vector concat_rhs = make_pattern_f32(32 * 7, 83, 0.0234375f); + set_tensor_pair_bytes(cpu_backend, cpu_concat_lhs, hrx_backend, hrx_concat_lhs, concat_lhs.data(), + concat_lhs.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_concat_rhs, hrx_backend, hrx_concat_rhs, concat_rhs.data(), + concat_rhs.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_unary), get_f32_tensor(cpu_backend, cpu_unary), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_copy), get_f32_tensor(cpu_backend, cpu_copy), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_concat), get_f32_tensor(cpu_backend, cpu_concat), 0.0f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + static void run_view_mul_f32_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -3134,6 +3783,15 @@ static void run_get_rows_f32_cpu_reference_case(ggml_type weight_type, ggml_build_forward_expand(hrx_graph, hrx_output); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_get_rows_f32" }); + const ggml::hrx::KernelSpecialization specialization = + scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_get_rows_f32"); + REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + int64_t expected_token_capacity = 1; + while (expected_token_capacity < token_count) { + expected_token_capacity *= 2; + } + REQUIRE(specialization.compile_parameters.at("ggml.get_rows_f32.token_capacity") == + std::to_string(expected_token_capacity)); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -3350,10 +4008,10 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const bool packed_input = packed_q4_input || packed_iq4_xs_input || square_output; const bool normalized_packed_input = normalize_input && packed_input; if (normalized_packed_input) { - require_kernel_subsequence(sequence, { "loom_libs:ggml_rmsnorm_binary_q8_1_x4", expected_kernel }); + require_kernel_subsequence(sequence, { "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish", expected_kernel }); REQUIRE(std::find(sequence.begin(), sequence.end(), "qwen3_moe:ggml_quantize_q8_1_x4_f32") == sequence.end()); } else if (normalize_input) { - REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:ggml_rmsnorm_binary_q8_1_x4") == sequence.end()); + REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish") == sequence.end()); } ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); @@ -3464,7 +4122,8 @@ static void run_prefill_value_cache_repeated_indices_case() { REQUIRE(setenv(disable_env, "1", 1) == 0); require_kernel_subsequence(scheduled_kernel_sequence(split_graph), { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", "loom_libs:ggml_set_rows" }); + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", + "loom_libs:ggml_set_rows" }); ggml_backend_buffer_t fused_buffer = ggml_backend_alloc_ctx_tensors(fused_ctx, fused_backend); ggml_backend_buffer_t split_buffer = ggml_backend_alloc_ctx_tensors(split_ctx, split_backend); @@ -4306,6 +4965,12 @@ static void run_rope_set_rows_cpu_reference_case() { ggml_build_forward_expand(cpu_graph, cpu_out); ggml_build_forward_expand(hrx_graph, hrx_out); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_f32" }); + const ggml::hrx::KernelSpecialization specialization = + scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_rope_f32"); + REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + REQUIRE(specialization.compile_parameters.at("ggml.rope_f32.token_capacity") == "4"); + REQUIRE(specialization.compile_parameters.count("ggml.rope_f32.input_span") == 0); + REQUIRE(specialization.integer_parameters.at("input_span") == hidden_size * token_count); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -5044,10 +5709,14 @@ static void run_decode_routed_moe_scheduling_case(ggml_type down_weight_type, bo REQUIRE(next_rms != nullptr); ggml_tensor * output = ggml_mul(ctx, next_rms, next_norm_weight); REQUIRE(output != nullptr); + ggml_tensor * next_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenHiddenSize); + ggml_tensor * next_output = ggml_mul_mat(ctx, next_weight, output); + REQUIRE(next_weight != nullptr); + REQUIRE(next_output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, next_output); std::vector expected = { "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", @@ -5242,7 +5911,9 @@ using test_runner::Suite; static void register_basic_ops_cases(Suite & suite) { suite.host_case("support.rmsnorm", [] { run_rmsnorm_support_checks(); }); suite.host_case("support.alternate_value_alias_lookup", [] { run_alternate_value_alias_lookup_checks(); }); + suite.host_case("support.activation_publication_contract", [] { run_activation_publication_contract_checks(); }); suite.device_case("add_f32.cpu_reference", [] { run_add_f32_cpu_reference_case(); }); + suite.device_case("dynamic_generic_ops.cpu_reference", [] { run_dynamic_generic_ops_cpu_reference_case(); }); for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".plain", [type] { run_ssm_conv_prefill_recurrent_case(type, false, false); }); @@ -5285,6 +5956,21 @@ static void register_basic_ops_cases(Suite & suite) { } static void register_dense_matmul_cases(Suite & suite) { + for (const ggml_type type : { GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M }) { + const char * type_suffix = type == GGML_TYPE_IQ1_S ? "iq1_s" : "iq1_m"; + const char * vector_kernel = type == GGML_TYPE_IQ1_S ? + "loom_libs:ggml_mul_mat_vector_iq1_s_f32_f32" : + "loom_libs:ggml_mul_mat_vector_iq1_m_f32_f32"; + const char * tiled_kernel = type == GGML_TYPE_IQ1_S ? + "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32" : + "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32"; + suite.device_case(std::string("dense_matmul.") + type_suffix + ".vector.tokens1.outputs128", [=] { + run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, 128, 2048); + }); + suite.device_case(std::string("dense_matmul.") + type_suffix + ".tiled.tokens64.outputs128", [=] { + run_dense_matmul_cpu_reference_case(type, tiled_kernel, 64, 128, 2048); + }); + } for (const int64_t tokens : { 1, 3, 5 }) { for (const int64_t outputs : { 1, 4, 47, 48, 63 }) { suite.device_case( @@ -5323,9 +6009,48 @@ static void register_dense_matmul_cases(Suite & suite) { suite.device_case("dense_matmul.q4_k.vector.tokens1.outputs128", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); }); + for (const ggml_type type : { GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS }) { + suite.device_case("dense_matmul." + type_name(type) + ".codebook.vector.tokens1.outputs128", [type] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32", 1, 128, + 2048); + }); + suite.device_case("dense_matmul." + type_name(type) + ".codebook.tiled.tokens64.outputs128", [type] { + run_dense_matmul_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", 64, 128, 2048); + }); + } + suite.device_case("dense_matmul.iq3_xxs.codebook.vector.tokens1.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ3_XXS, + "loom_libs:ggml_mul_mat_vector_iq3_xxs_f32_f32", 1, 128, 2048); + }); + suite.device_case("dense_matmul.iq3_xxs.codebook.tiled.tokens64.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ3_XXS, + "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", 64, + 128, 2048); + }); + for (const ggml_type type : { GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_S }) { + const char * vector_kernel = type == GGML_TYPE_IQ2_S ? + "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32" : + "loom_libs:ggml_mul_mat_vector_iq3_s_f32_f32"; + const char * tiled_kernel = type == GGML_TYPE_IQ2_S ? + "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32" : + "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32"; + suite.device_case("dense_matmul." + type_name(type) + ".codebook_control.vector.tokens1.outputs128", + [type, vector_kernel] { + run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, 128, 2048); + }); + suite.device_case("dense_matmul." + type_name(type) + ".codebook_control.tiled.tokens64.outputs128", + [type, tiled_kernel] { + run_dense_matmul_cpu_reference_case(type, tiled_kernel, 64, 128, 2048); + }); + } suite.device_case("dense_matmul.iq4_xs.vector.tokens1.outputs128", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); }); + suite.device_case("dense_matmul.iq4_xs.tiled.tokens64.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", 64, 128); + }); suite.device_case("dense_matmul.iq4_xs.prefill.tokens256.outputs128", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 128, @@ -5366,6 +6091,11 @@ static void register_dense_matmul_cases(Suite & suite) { run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); }); } + suite.device_case("dense_matmul.iq4_nl.tiled.tokens64.outputs1024", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", 64, 1024, + 640); + }); suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs256.input640", [] { run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, 256, 640); @@ -5408,6 +6138,12 @@ static void register_dense_matmul_cases(Suite & suite) { "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 2048, GGML_GLU_OP_GEGLU, 640); }); + suite.device_case("dense_matmul_swiglu.iq4_nl.iq4_nl.geglu.tokens64.outputs128.input640", [] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", + 64, 128, GGML_GLU_OP_GEGLU, 640); + }); for (const ggml_glu_op op : { GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { suite.device_case(std::string("dense_matmul_swiglu.q4_k.f16.") + glu_op_name(op) + ".tokens2.outputs128", [op] { run_dense_matmul_swiglu_cpu_reference_case( @@ -5684,6 +6420,10 @@ static void register_rmsnorm_and_scheduling_cases(Suite & suite) { } static void register_flash_attention_cases(Suite & suite) { + 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"); + }); 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", [head_size] {